channel.py 33.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# pylint: disable=doc-string-missing
B
barriery 已提交
15
from time import time as _time
D
dongdaxiang 已提交
16 17 18 19 20 21 22 23 24 25
import threading
import multiprocessing
import multiprocessing.queues
import sys
if sys.version_info.major == 2:
    import Queue
elif sys.version_info.major == 3:
    import queue as Queue
else:
    raise Exception("Error Python version")
26 27 28
import numpy as np
import logging
import enum
29
import os
30
import copy
D
dongdaxiang 已提交
31

32
_LOGGER = logging.getLogger(__name__)
B
barrierye 已提交
33

D
dongdaxiang 已提交
34

T
TeslaZhao 已提交
35 36 37 38
class ChannelDataErrcode(enum.Enum):
    """
    ChannelData error code
    """
D
dongdaxiang 已提交
39 40 41 42 43
    OK = 0
    TIMEOUT = 1
    NOT_IMPLEMENTED = 2
    TYPE_ERROR = 3
    RPC_PACKAGE_ERROR = 4
B
barrierye 已提交
44
    CLIENT_ERROR = 5
B
barrierye 已提交
45
    CLOSED_ERROR = 6
B
barriery 已提交
46 47
    NO_SERVICE = 7
    UNKNOW = 8
T
TeslaZhao 已提交
48 49 50 51 52 53 54 55 56
    PRODUCT_ERROR = 9


class ProductErrCode(enum.Enum):
    """
    ProductErrCode is a base class for recording business error code. 
    product developers inherit this class and extend more error codes. 
    """
    pass
D
dongdaxiang 已提交
57 58 59 60 61 62 63 64


class ChannelDataType(enum.Enum):
    DICT = 0
    CHANNEL_NPDATA = 1
    ERROR = 2


65 66 67 68
class ChannelData(object):
    def __init__(self,
                 datatype=None,
                 npdata=None,
B
barrierye 已提交
69
                 dictdata=None,
70
                 data_id=None,
T
TeslaZhao 已提交
71 72
                 log_id=None,
                 error_code=None,
B
barrierye 已提交
73
                 error_info=None,
T
TeslaZhao 已提交
74 75
                 prod_error_code=None,
                 prod_error_info=None,
B
barrierye 已提交
76
                 client_need_profile=False):
77 78 79
        '''
        There are several ways to use it:
        
T
TeslaZhao 已提交
80 81 82
        1. ChannelData(ChannelDataType.CHANNEL_NPDATA.value, npdata, data_id, log_id)
        2. ChannelData(ChannelDataType.DICT.value, dictdata, data_id, log_id)
        3. ChannelData(error_code, error_info, prod_error_code, prod_error_info, data_id, log_id)
83 84 85 86

        Protobufs are not pickle-able:
        https://stackoverflow.com/questions/55344376/how-to-import-protobuf-module
        '''
T
TeslaZhao 已提交
87
        if error_code is not None or prod_error_code is not None:
88
            if data_id is None or error_info is None:
B
barriery 已提交
89 90
                _LOGGER.critical("Failed to generate ChannelData: data_id"
                                 " and error_info cannot be None")
91
                os._exit(-1)
92 93
            datatype = ChannelDataType.ERROR.value
        else:
B
barrierye 已提交
94
            if datatype == ChannelDataType.CHANNEL_NPDATA.value:
T
TeslaZhao 已提交
95 96
                error_code, error_info = ChannelData.check_npdata(npdata)
                if error_code != ChannelDataErrcode.OK.value:
B
barrierye 已提交
97
                    datatype = ChannelDataType.ERROR.value
T
TeslaZhao 已提交
98 99
                    _LOGGER.error("(data_id={} log_id={}) {}".format(
                        data_id, log_id, error_info))
B
barrierye 已提交
100
            elif datatype == ChannelDataType.DICT.value:
T
TeslaZhao 已提交
101 102
                error_code, error_info = ChannelData.check_dictdata(dictdata)
                if error_code != ChannelDataErrcode.OK.value:
B
barrierye 已提交
103
                    datatype = ChannelDataType.ERROR.value
T
TeslaZhao 已提交
104 105
                    _LOGGER.error("(data_id={} log_id={}) {}".format(
                        data_id, log_id, error_info))
106
            else:
T
TeslaZhao 已提交
107 108
                _LOGGER.critical("(data_id={} log_id={}) datatype not match".
                                 format(data_id, log_id))
109
                os._exit(-1)
110
        self.datatype = datatype
B
barrierye 已提交
111 112
        self.npdata = npdata
        self.dictdata = dictdata
113
        self.id = data_id
T
TeslaZhao 已提交
114 115
        self.log_id = log_id
        self.error_code = error_code
116
        self.error_info = error_info
T
TeslaZhao 已提交
117 118
        self.prod_error_code = prod_error_code
        self.prod_error_info = prod_error_info
B
barrierye 已提交
119
        self.client_need_profile = client_need_profile
B
barrierye 已提交
120
        self.profile_data_set = set()
B
barrierye 已提交
121

B
barrierye 已提交
122
    def add_profile(self, profile_set):
B
barrierye 已提交
123 124
        if self.client_need_profile is False:
            self.client_need_profile = True
B
barrierye 已提交
125
        self.profile_data_set |= profile_set
126

B
barrierye 已提交
127 128
    @staticmethod
    def check_dictdata(dictdata):
T
TeslaZhao 已提交
129
        error_code = ChannelDataErrcode.OK.value
B
barrierye 已提交
130
        error_info = None
B
barrierye 已提交
131 132 133 134
        if isinstance(dictdata, list):
            # batch data
            for sample in dictdata:
                if not isinstance(sample, dict):
T
TeslaZhao 已提交
135
                    error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
136 137
                    error_info = "Failed to check data: the type of " \
                            "data must be dict, but get {}.".format(type(sample))
B
barrierye 已提交
138 139 140
                    break
        elif not isinstance(dictdata, dict):
            # batch size = 1
T
TeslaZhao 已提交
141
            error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
142 143
            error_info = "Failed to check data: the type of data must " \
                    "be dict, but get {}.".format(type(dictdata))
T
TeslaZhao 已提交
144
        return error_code, error_info
B
barrierye 已提交
145

B
bug fix  
barriery 已提交
146 147
    @staticmethod
    def check_batch_npdata(batch):
T
TeslaZhao 已提交
148
        error_code = ChannelDataErrcode.OK.value
B
bug fix  
barriery 已提交
149 150
        error_info = None
        for npdata in batch:
T
TeslaZhao 已提交
151 152
            error_code, error_info = ChannelData.check_npdata(npdata)
            if error_code != ChannelDataErrcode.OK.value:
B
bug fix  
barriery 已提交
153
                break
T
TeslaZhao 已提交
154
        return error_code, error_info
B
bug fix  
barriery 已提交
155

B
barrierye 已提交
156 157
    @staticmethod
    def check_npdata(npdata):
T
TeslaZhao 已提交
158
        error_code = ChannelDataErrcode.OK.value
159
        error_info = None
W
wangjiawei04 已提交
160 161 162 163
        if isinstance(npdata, list):
            # batch data
            for sample in npdata:
                if not isinstance(sample, dict):
T
TeslaZhao 已提交
164
                    error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
165 166 167
                    error_info = "Failed to check data: the " \
                            "value of data must be dict, but get {}.".format(
                                    type(sample))
W
wangjiawei04 已提交
168 169 170
                    break
                for _, value in sample.items():
                    if not isinstance(value, np.ndarray):
T
TeslaZhao 已提交
171
                        error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
172 173 174
                        error_info = "Failed to check data: the" \
                                " value of data must be np.ndarray, but get {}.".format(
                                        type(value))
T
TeslaZhao 已提交
175
                        return error_code, error_info
W
wangjiawei04 已提交
176 177 178 179
        elif isinstance(npdata, dict):
            # batch_size = 1
            for _, value in npdata.items():
                if not isinstance(value, np.ndarray):
T
TeslaZhao 已提交
180
                    error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
181 182 183
                    error_info = "Failed to check data: the value " \
                            "of data must be np.ndarray, but get {}.".format(
                                    type(value))
W
wangjiawei04 已提交
184 185
                    break
        else:
T
TeslaZhao 已提交
186
            error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
187 188
            error_info = "Failed to check data: the value of data " \
                    "must be dict, but get {}.".format(type(npdata))
T
TeslaZhao 已提交
189
        return error_code, error_info
190 191 192

    def parse(self):
        feed = None
B
barrierye 已提交
193 194
        if self.datatype == ChannelDataType.CHANNEL_NPDATA.value:
            # return narray
195
            feed = self.npdata
B
barrierye 已提交
196 197 198
        elif self.datatype == ChannelDataType.DICT.value:
            # return dict
            feed = self.dictdata
199
        else:
B
barriery 已提交
200 201
            _LOGGER.critical("Failed to parse channeldata: error " \
                    "type({}) in datatype.".format(self.datatype))
202
            os._exit(-1)
203 204
        return feed

205 206 207 208 209 210 211 212
    def __cmp__(self, other):
        if self.id < other.id:
            return -1
        elif self.id == other.id:
            return 0
        else:
            return 1

213
    def __str__(self):
T
TeslaZhao 已提交
214 215 216
        return "type[{}], error_code[{}], data_id[{}], log_id[{}]".format(
            ChannelDataType(self.datatype).name, self.error_code, self.id,
            self.log_id)
217 218


B
barrierye 已提交
219
class ProcessChannel(object):
220 221 222 223 224 225 226 227 228
    """ 
    (Process version) The channel used for communication between Ops.

    1. Support multiple different Op feed data (multiple producer)
        Different types of data will be packaged through the data ID
    2. Support multiple different Op fetch data (multiple consumer)
        Only when all types of Ops get the data of the same ID,
        the data will be poped; The Op of the same type will not
        get the data of the same ID.
B
barriery 已提交
229
    3. Function front support timeout param to make auto-batching.
230 231 232 233 234

    Note:
    1. The ID of the data in the channel must be different.
    2. The function add_producer() and add_consumer() are not thread safe,
       and can only be called during initialization.
B
barrierye 已提交
235 236 237 238 239 240 241 242 243 244 245

    There are two buffers and one queue in Channel:

        op_A \                                           / op_D
        op_B - a. input_buf -> b. queue -> c. output_buf - op_E
        op_C /                                           \ op_F
    
    a. In input_buf, the input of multiple predecessor Ops is packed by data ID.
    b. The packed data will be stored in queue.
    c. In order to support multiple successor Ops to retrieve data, output_buf
        maintains the data obtained from queue.
246 247
    """

B
barriery 已提交
248
    def __init__(self, manager, name=None, maxsize=0):
B
barrierye 已提交
249 250 251 252 253 254
        # For queue multiprocess: after putting an object on 
        # an empty queue there may be an infinitessimal delay
        # before the queue's :meth:`~Queue.empty`
        # see more:
        # - https://bugs.python.org/issue18277
        # - https://hg.python.org/cpython/rev/860fc6a2bd21
255
        self._que = manager.PriorityQueue(maxsize=maxsize)
256 257
        self._maxsize = maxsize
        self.name = name
258
        self._stop = manager.Value('i', 0)
259 260 261 262

        self._cv = multiprocessing.Condition()

        self._producers = []
B
barrierye 已提交
263
        self._pushed_producer_count = manager.dict()  # {data_id: count}
B
barrierye 已提交
264
        self._input_buf = manager.dict()  # {data_id: {op_name: data}}
265

266
        self._reset_max_cursor = 1000000000000000000
B
barrierye 已提交
267 268 269 270
        self._consumer_cursors = manager.dict()  # {op_name: cursor}
        self._cursor_count = manager.dict()  # {cursor: count}
        self._base_cursor = manager.Value('i', 0)
        self._output_buf = manager.list()
271

B
barriery 已提交
272 273 274
    def get_maxsize(self):
        return self._maxsize

B
barriery 已提交
275 276 277
    def size(self):
        return self._que.qsize()

278 279 280 281
    def get_producers(self):
        return self._producers

    def get_consumers(self):
B
barrierye 已提交
282
        return self._consumer_cursors.keys()
283 284 285 286 287 288 289

    def _log(self, info_str):
        return "[{}] {}".format(self.name, info_str)

    def add_producer(self, op_name):
        """ not thread safe, and can only be called during initialization. """
        if op_name in self._producers:
290
            _LOGGER.critical(
B
barriery 已提交
291 292
                self._log("Failed to add producer: producer({})" \
                        " is already in channel".format(op_name)))
293
            os._exit(-1)
294
        self._producers.append(op_name)
B
barriery 已提交
295
        _LOGGER.debug(self._log("Succ add a producer: {}".format(op_name)))
296 297 298

    def add_consumer(self, op_name):
        """ not thread safe, and can only be called during initialization. """
B
barrierye 已提交
299
        if op_name in self._consumer_cursors:
300
            _LOGGER.critical(
B
barriery 已提交
301 302
                    self._log("Failed to add consumer: consumer({})" \
                            " is already in channel".format(op_name)))
303
            os._exit(-1)
B
barrierye 已提交
304
        self._consumer_cursors[op_name] = 0
305

B
barrierye 已提交
306 307 308
        if self._cursor_count.get(0) is None:
            self._cursor_count[0] = 0
        self._cursor_count[0] += 1
B
barriery 已提交
309
        _LOGGER.debug(self._log("Succ add a consumer: {}".format(op_name)))
310 311

    def push(self, channeldata, op_name=None):
B
barrierye 已提交
312
        _LOGGER.debug(
T
TeslaZhao 已提交
313 314
            self._log("(data_id={} log_id={}) Op({}) Pushing data".format(
                channeldata.id, channeldata.log_id, op_name)))
315
        if len(self._producers) == 0:
316
            _LOGGER.critical(
317
                self._log(
T
TeslaZhao 已提交
318
                    "(data_id={} log_id={}) Op({}) Failed to push data: expected number"
B
barriery 已提交
319
                    " of producers to be greater than 0, but the it is 0.".
T
TeslaZhao 已提交
320
                    format(channeldata.id, channeldata.log_id, op_name)))
321
            os._exit(-1)
322 323
        elif len(self._producers) == 1:
            with self._cv:
324
                while self._stop.value == 0:
325
                    try:
B
barrierye 已提交
326
                        self._que.put({op_name: channeldata}, timeout=0)
327 328 329
                        break
                    except Queue.Full:
                        self._cv.wait()
330
                if self._stop.value == 1:
B
barrierye 已提交
331
                    raise ChannelStopError()
332
                self._cv.notify_all()
B
barriery 已提交
333
            _LOGGER.debug(
T
TeslaZhao 已提交
334 335 336
                self._log(
                    "(data_id={} log_id={}) Op({}) Pushed data into internal queue.".
                    format(channeldata.id, channeldata.log_id, op_name)))
337 338
            return True
        elif op_name is None:
339
            _LOGGER.critical(
340
                self._log(
T
TeslaZhao 已提交
341
                    "(data_id={} log_id={}) Op({}) Failed to push data: there are multiple "
B
barriery 已提交
342
                    "producers, so op_name cannot be None.".format(
T
TeslaZhao 已提交
343
                        channeldata.id, channeldata.log_id, op_name)))
344
            os._exit(-1)
345 346 347

        producer_num = len(self._producers)
        data_id = channeldata.id
T
TeslaZhao 已提交
348
        log_id = channeldata.log_id
349 350
        put_data = None
        with self._cv:
B
barrierye 已提交
351 352
            if data_id not in self._input_buf:
                self._input_buf[data_id] = {
353 354 355
                    name: None
                    for name in self._producers
                }
B
barrierye 已提交
356
                self._pushed_producer_count[data_id] = 0
357
            # see: https://docs.python.org/3.6/library/multiprocessing.html?highlight=multiprocess#proxy-objects
B
barrierye 已提交
358 359 360 361 362
            # self._input_buf[data_id][op_name] = channeldata
            tmp_input_buf = self._input_buf[data_id]
            tmp_input_buf[op_name] = channeldata
            self._input_buf[data_id] = tmp_input_buf

B
barrierye 已提交
363
            if self._pushed_producer_count[data_id] + 1 == producer_num:
B
barrierye 已提交
364 365
                put_data = self._input_buf[data_id]
                self._input_buf.pop(data_id)
B
barrierye 已提交
366
                self._pushed_producer_count.pop(data_id)
367
            else:
B
barrierye 已提交
368
                self._pushed_producer_count[data_id] += 1
369 370

            if put_data is None:
B
barrierye 已提交
371
                _LOGGER.debug(
B
barriery 已提交
372
                    self._log(
T
TeslaZhao 已提交
373 374
                        "(data_id={} log_id={}) Op({}) Pushed data into input_buffer.".
                        format(data_id, log_id, op_name)))
375
            else:
376
                while self._stop.value == 0:
377
                    try:
B
barrierye 已提交
378
                        self._que.put(put_data, timeout=0)
379 380 381
                        break
                    except Queue.Empty:
                        self._cv.wait()
382
                if self._stop.value == 1:
B
barrierye 已提交
383
                    raise ChannelStopError()
384

B
barrierye 已提交
385
                _LOGGER.debug(
B
barriery 已提交
386
                    self._log(
T
TeslaZhao 已提交
387 388
                        "(data_id={} log_id={}) Op({}) Pushed data into internal_queue.".
                        format(data_id, log_id, op_name)))
389 390 391
            self._cv.notify_all()
        return True

B
barriery 已提交
392
    def front(self, op_name=None, timeout=None):
B
barriery 已提交
393
        _LOGGER.debug(
B
barriery 已提交
394 395
            self._log("Op({}) Getting data[?]; timeout(s)={}".format(op_name,
                                                                     timeout)))
B
barriery 已提交
396
        endtime = None
B
bug fix  
barriery 已提交
397 398 399 400 401
        if timeout is not None:
            if timeout <= 0:
                timeout = None
            else:
                endtime = _time() + timeout
B
barriery 已提交
402

B
barrierye 已提交
403
        if len(self._consumer_cursors) == 0:
404
            _LOGGER.critical(
405
                self._log(
B
barriery 已提交
406 407
                    "Op({}) Failed to get data: expected number of consumers to be " \
                            "greater than 0, but the it is 0.".format(op_name)))
408
            os._exit(-1)
B
barrierye 已提交
409
        elif len(self._consumer_cursors) == 1:
410 411
            resp = None
            with self._cv:
412
                while self._stop.value == 0 and resp is None:
413
                    try:
B
barrierye 已提交
414
                        resp = self._que.get(timeout=0)
415 416
                        break
                    except Queue.Empty:
B
barriery 已提交
417 418 419
                        if timeout is not None:
                            remaining = endtime - _time()
                            if remaining <= 0.0:
B
barriery 已提交
420
                                _LOGGER.debug(
B
barriery 已提交
421 422
                                    self._log("Op({}) Failed to get data: "
                                              "timeout".format(op_name)))
B
barriery 已提交
423 424 425 426
                                raise ChannelTimeoutError()
                            self._cv.wait(remaining)
                        else:
                            self._cv.wait()
427
                if self._stop.value == 1:
B
barrierye 已提交
428
                    raise ChannelStopError()
B
barriery 已提交
429
            _LOGGER.debug(
T
TeslaZhao 已提交
430 431
                self._log("(data_id={} log_id={}) Op({}) Got data".format(
                    resp.values()[0].id, resp.values()[0].log_id, op_name)))
432 433
            return resp
        elif op_name is None:
434
            _LOGGER.critical(
435
                self._log(
B
barriery 已提交
436 437
                    "Op({}) Failed to get data: there are multiple consumers, "
                    "so op_name cannot be None.".format(op_name)))
438
            os._exit(-1)
439

B
barrierye 已提交
440 441 442 443 444 445 446 447 448 449
        # In output_buf, different Ops (according to op_name) have different
        # cursors. In addition, there is a base_cursor. Their difference is
        # the data_idx to be taken by the corresponding Op at the current
        # time:    data_idx = consumer_cursor - base_cursor
        # 
        #            base_cursor    consumer_B_cursor (data_idx: 3)
        #                 |                       |
        # output_buf: | data0 | data1 | data2 | data3 |
        #                 |
        #   consumer_A_cursor (data_idx: 0)
450
        with self._cv:
B
barrierye 已提交
451 452
            # When the data required by the current Op is not in output_buf,
            # it is necessary to obtain a data from queue and add it to output_buf.
453
            while self._stop.value == 0 and self._consumer_cursors[
B
barrierye 已提交
454
                    op_name] - self._base_cursor.value >= len(self._output_buf):
455
                try:
B
barrierye 已提交
456
                    channeldata = self._que.get(timeout=0)
B
barrierye 已提交
457
                    self._output_buf.append(channeldata)
B
barriery 已提交
458
                    _LOGGER.debug(
B
barriery 已提交
459
                        self._log(
T
TeslaZhao 已提交
460 461 462
                            "(data_id={} log_id={}) Op({}) Pop ready item into output_buffer".
                            format(channeldata.values()[0].id,
                                   channeldata.values()[0].log_id, op_name)))
463 464
                    break
                except Queue.Empty:
B
barriery 已提交
465 466 467
                    if timeout is not None:
                        remaining = endtime - _time()
                        if remaining <= 0.0:
B
barriery 已提交
468
                            _LOGGER.debug(
B
barriery 已提交
469 470
                                self._log("Op({}) Failed to get data: timeout".
                                          format(op_name)))
B
barriery 已提交
471 472 473 474
                            raise ChannelTimeoutError()
                        self._cv.wait(remaining)
                    else:
                        self._cv.wait()
475
            if self._stop.value == 1:
B
barrierye 已提交
476
                raise ChannelStopError()
477

B
barrierye 已提交
478 479 480 481
            consumer_cursor = self._consumer_cursors[op_name]
            base_cursor = self._base_cursor.value
            data_idx = consumer_cursor - base_cursor
            resp = self._output_buf[data_idx]
482

B
barrierye 已提交
483 484 485 486 487 488 489 490
            self._cursor_count[consumer_cursor] -= 1
            if consumer_cursor == base_cursor and self._cursor_count[
                    consumer_cursor] == 0:
                # When all the different Ops get the data that data_idx points
                # to, pop the data from output_buf.
                self._cursor_count.pop(consumer_cursor)
                self._output_buf.pop(0)
                self._base_cursor.value += 1
491 492
                # to avoid cursor overflow
                if self._base_cursor.value >= self._reset_max_cursor:
B
barriery 已提交
493
                    _LOGGER.info(self._log("Reset cursor in Channel"))
494 495 496 497 498 499 500 501 502 503
                    self._base_cursor.value -= self._reset_max_cursor
                    for name in self._consumer_cursors.keys():
                        self._consumer_cursors[name] -= self._reset_max_cursor
                    cursor_count_tmp = {
                        cursor - self._reset_max_cursor: count
                        for cursor, count in self._cursor_count.copy().items()
                    }
                    self._cursor_count.clear()
                    for cursor, count in cursor_count_tmp.items():
                        self._cursor_count[cursor] = count
B
barrierye 已提交
504 505 506 507 508 509

            self._consumer_cursors[op_name] += 1
            new_consumer_cursor = self._consumer_cursors[op_name]
            if self._cursor_count.get(new_consumer_cursor) is None:
                self._cursor_count[new_consumer_cursor] = 0
            self._cursor_count[new_consumer_cursor] += 1
510

511 512
            self._cv.notify_all()

B
barriery 已提交
513
        _LOGGER.debug(
T
TeslaZhao 已提交
514 515 516
            self._log(
                "(data_id={} log_id={}) Op({}) Got data from output_buffer".
                format(resp.values()[0].id, resp.values()[0].log_id, op_name)))
B
barriery 已提交
517
        return resp
518 519

    def stop(self):
520
        _LOGGER.info(self._log("stop."))
521
        self._stop.value = 1
B
barrierye 已提交
522 523
        with self._cv:
            self._cv.notify_all()
524 525


526
class ThreadChannel(Queue.PriorityQueue):
527 528 529 530 531 532 533 534 535
    """ 
    (Thread version)The channel used for communication between Ops.

    1. Support multiple different Op feed data (multiple producer)
        Different types of data will be packaged through the data ID
    2. Support multiple different Op fetch data (multiple consumer)
        Only when all types of Ops get the data of the same ID,
        the data will be poped; The Op of the same type will not
        get the data of the same ID.
B
barriery 已提交
536
    3. Function front support timeout param to make auto-batching.
537 538 539 540 541

    Note:
    1. The ID of the data in the channel must be different.
    2. The function add_producer() and add_consumer() are not thread safe,
       and can only be called during initialization.
B
barrierye 已提交
542 543 544 545 546 547 548 549 550 551 552

    There are two buffers and one queue in Channel:

        op_A \                                           / op_D
        op_B - a. input_buf -> b. queue -> c. output_buf - op_E
        op_C /                                           \ op_F
    
    a. In input_buf, the input of multiple predecessor Ops is packed by data ID.
    b. The packed data will be stored in queue.
    c. In order to support multiple successor Ops to retrieve data, output_buf
        maintains the data obtained from queue.
553 554
    """

B
barriery 已提交
555
    def __init__(self, name=None, maxsize=-1):
556 557 558 559 560 561 562 563
        Queue.Queue.__init__(self, maxsize=maxsize)
        self._maxsize = maxsize
        self.name = name
        self._stop = False

        self._cv = threading.Condition()

        self._producers = []
B
barrierye 已提交
564
        self._pushed_producer_count = {}  # {data_id: count}
B
barrierye 已提交
565
        self._input_buf = {}  # {data_id: {op_name: data}}
566

567
        self._reset_max_cursor = 1000000000000000000
B
barrierye 已提交
568 569 570 571
        self._consumer_cursors = {}  # {op_name: idx}
        self._cursor_count = {}  # {cursor: count}
        self._base_cursor = 0
        self._output_buf = []
572

B
barriery 已提交
573 574 575
    def get_maxsize(self):
        return self._maxsize

B
barriery 已提交
576 577 578
    def size(self):
        return self.qsize()

579 580 581 582
    def get_producers(self):
        return self._producers

    def get_consumers(self):
B
barrierye 已提交
583
        return self._consumer_cursors.keys()
584 585 586 587 588 589 590

    def _log(self, info_str):
        return "[{}] {}".format(self.name, info_str)

    def add_producer(self, op_name):
        """ not thread safe, and can only be called during initialization. """
        if op_name in self._producers:
591
            _LOGGER.critical(
B
barriery 已提交
592 593
                self._log("Failed to add producer: producer({}) is "
                          "already in channel".format(op_name)))
594
            os._exit(-1)
595
        self._producers.append(op_name)
B
barriery 已提交
596
        _LOGGER.debug(self._log("Succ add a producer: {}".format(op_name)))
597 598 599

    def add_consumer(self, op_name):
        """ not thread safe, and can only be called during initialization. """
B
barrierye 已提交
600
        if op_name in self._consumer_cursors:
601
            _LOGGER.critical(
B
barriery 已提交
602 603
                self._log("Failed to add consumer: consumer({}) is "
                          "already in channel".format(op_name)))
604
            os._exit(-1)
B
barrierye 已提交
605
        self._consumer_cursors[op_name] = 0
606

B
barrierye 已提交
607 608 609
        if self._cursor_count.get(0) is None:
            self._cursor_count[0] = 0
        self._cursor_count[0] += 1
B
barriery 已提交
610
        _LOGGER.debug(self._log("Succ add a consumer: {}".format(op_name)))
611 612

    def push(self, channeldata, op_name=None):
B
barrierye 已提交
613
        _LOGGER.debug(
T
TeslaZhao 已提交
614 615
            self._log("(data_id={} log_id={}) Op({}) Pushing data".format(
                channeldata.id, channeldata.log_id, op_name)))
616
        if len(self._producers) == 0:
617
            _LOGGER.critical(
618
                self._log(
T
TeslaZhao 已提交
619
                    "(data_id={} log_id={}) Op({}) Failed to push data: expected number of "
B
barriery 已提交
620
                    "producers to be greater than 0, but the it is 0.".format(
T
TeslaZhao 已提交
621
                        channeldata.id, channeldata.log_id, op_name)))
622
            os._exit(-1)
623 624 625 626
        elif len(self._producers) == 1:
            with self._cv:
                while self._stop is False:
                    try:
B
barrierye 已提交
627
                        self.put({op_name: channeldata}, timeout=0)
628 629 630
                        break
                    except Queue.Full:
                        self._cv.wait()
B
barrierye 已提交
631 632
                if self._stop:
                    raise ChannelStopError()
633
                self._cv.notify_all()
B
barriery 已提交
634
            _LOGGER.debug(
T
TeslaZhao 已提交
635 636 637
                self._log(
                    "(data_id={} log_id={}) Op({}) Pushed data into internal_queue.".
                    format(channeldata.id, channeldata.log_id, op_name)))
638 639
            return True
        elif op_name is None:
640
            _LOGGER.critical(
641
                self._log(
T
TeslaZhao 已提交
642
                    "(data_id={} log_id={}) Op({}) Failed to push data: there are multiple"
B
barriery 已提交
643
                    " producers, so op_name cannot be None.".format(
T
TeslaZhao 已提交
644
                        channeldata.id, channeldata.log_id, op_name)))
645
            os._exit(-1)
646 647 648

        producer_num = len(self._producers)
        data_id = channeldata.id
T
TeslaZhao 已提交
649
        log_id = channeldata.log_id
650 651
        put_data = None
        with self._cv:
B
barrierye 已提交
652 653
            if data_id not in self._input_buf:
                self._input_buf[data_id] = {
654 655 656
                    name: None
                    for name in self._producers
                }
B
barrierye 已提交
657
                self._pushed_producer_count[data_id] = 0
B
barrierye 已提交
658
            self._input_buf[data_id][op_name] = channeldata
B
barrierye 已提交
659
            if self._pushed_producer_count[data_id] + 1 == producer_num:
B
barrierye 已提交
660 661
                put_data = self._input_buf[data_id]
                self._input_buf.pop(data_id)
B
barrierye 已提交
662
                self._pushed_producer_count.pop(data_id)
663
            else:
B
barrierye 已提交
664
                self._pushed_producer_count[data_id] += 1
665 666

            if put_data is None:
B
barrierye 已提交
667
                _LOGGER.debug(
B
barriery 已提交
668
                    self._log(
T
TeslaZhao 已提交
669 670
                        "(data_id={} log_id={}) Op({}) Pushed data into input_buffer.".
                        format(data_id, log_id, op_name)))
671 672 673 674 675 676 677
            else:
                while self._stop is False:
                    try:
                        self.put(put_data, timeout=0)
                        break
                    except Queue.Empty:
                        self._cv.wait()
B
barrierye 已提交
678 679
                if self._stop:
                    raise ChannelStopError()
680

B
barrierye 已提交
681
                _LOGGER.debug(
B
barriery 已提交
682
                    self._log(
T
TeslaZhao 已提交
683 684
                        "(data_id={} log_id={}) Op({}) Pushed data into internal_queue.".
                        format(data_id, log_id, op_name)))
685 686 687
            self._cv.notify_all()
        return True

B
barriery 已提交
688
    def front(self, op_name=None, timeout=None):
B
barriery 已提交
689
        _LOGGER.debug(
B
barriery 已提交
690 691
            self._log("Op({}) Getting data[?]; timeout(s)={}".format(op_name,
                                                                     timeout)))
B
barriery 已提交
692
        endtime = None
B
bug fix  
barriery 已提交
693 694 695 696 697
        if timeout is not None:
            if timeout <= 0:
                timeout = None
            else:
                endtime = _time() + timeout
B
barriery 已提交
698

B
barrierye 已提交
699
        if len(self._consumer_cursors) == 0:
700
            _LOGGER.critical(
701
                self._log(
B
barriery 已提交
702 703
                    "Op({}) Failed to get data: expected number of consumers to be "
                    "greater than 0, but the it is 0.".format(op_name)))
704
            os._exit(-1)
B
barrierye 已提交
705
        elif len(self._consumer_cursors) == 1:
706 707 708 709 710 711 712
            resp = None
            with self._cv:
                while self._stop is False and resp is None:
                    try:
                        resp = self.get(timeout=0)
                        break
                    except Queue.Empty:
B
barriery 已提交
713 714 715
                        if timeout is not None:
                            remaining = endtime - _time()
                            if remaining <= 0.0:
B
barriery 已提交
716
                                _LOGGER.debug(
B
barriery 已提交
717 718 719
                                    self._log(
                                        "Op({}) Failed to get data: timeout".
                                        format(op_name)))
B
barriery 已提交
720 721 722 723
                                raise ChannelTimeoutError()
                            self._cv.wait(remaining)
                        else:
                            self._cv.wait()
B
barrierye 已提交
724 725
                if self._stop:
                    raise ChannelStopError()
B
barrierye 已提交
726
            _LOGGER.debug(
T
TeslaZhao 已提交
727 728
                self._log("(data_id={} log_id={}) Op({}) Got data".format(
                    resp.values()[0].id, resp.values()[0].log_id, op_name)))
729 730
            return resp
        elif op_name is None:
731
            _LOGGER.critical(
B
barriery 已提交
732 733 734
                self._log("Op({}) Failed to get data: there are multiple "
                          "consumers, so op_name cannot be None.".format(
                              op_name)))
735
            os._exit(-1)
736

B
barrierye 已提交
737 738 739 740 741 742 743 744 745 746
        # In output_buf, different Ops (according to op_name) have different
        # cursors. In addition, there is a base_cursor. Their difference is
        # the data_idx to be taken by the corresponding Op at the current
        # time:    data_idx = consumer_cursor - base_cursor
        # 
        #            base_cursor    consumer_B_cursor (data_idx: 3)
        #                 |                       |
        # output_buf: | data0 | data1 | data2 | data3 |
        #                 |
        #   consumer_A_cursor (data_idx: 0)
747
        with self._cv:
B
barrierye 已提交
748 749 750 751
            # When the data required by the current Op is not in output_buf,
            # it is necessary to obtain a data from queue and add it to output_buf.
            while self._stop is False and self._consumer_cursors[
                    op_name] - self._base_cursor >= len(self._output_buf):
752 753
                try:
                    channeldata = self.get(timeout=0)
B
barrierye 已提交
754
                    self._output_buf.append(channeldata)
B
barriery 已提交
755
                    _LOGGER.debug(
B
barriery 已提交
756
                        self._log(
T
TeslaZhao 已提交
757 758 759
                            "(data_id={} log_id={}) Op({}) Pop ready item into output_buffer".
                            format(channeldata.values()[0].id,
                                   channeldata.values()[0].log_id, op_name)))
760 761
                    break
                except Queue.Empty:
B
barriery 已提交
762 763 764
                    if timeout is not None:
                        remaining = endtime - _time()
                        if remaining <= 0.0:
B
barriery 已提交
765
                            _LOGGER.debug(
B
barriery 已提交
766 767
                                self._log("Op({}) Failed to get data: timeout".
                                          format(op_name)))
B
barriery 已提交
768 769 770 771
                            raise ChannelTimeoutError()
                        self._cv.wait(remaining)
                    else:
                        self._cv.wait()
B
barrierye 已提交
772 773
            if self._stop:
                raise ChannelStopError()
774

B
barrierye 已提交
775 776 777
            consumer_cursor = self._consumer_cursors[op_name]
            base_cursor = self._base_cursor
            data_idx = consumer_cursor - base_cursor
B
barrierye 已提交
778 779

            resp = None
780

B
barrierye 已提交
781 782 783 784 785 786
            self._cursor_count[consumer_cursor] -= 1
            if consumer_cursor == base_cursor and self._cursor_count[
                    consumer_cursor] == 0:
                # When all the different Ops get the data that data_idx points
                # to, pop the data from output_buf.
                self._cursor_count.pop(consumer_cursor)
B
barrierye 已提交
787
                resp = self._output_buf.pop(0)
B
barrierye 已提交
788
                self._base_cursor += 1
789 790
                # to avoid cursor overflow
                if self._base_cursor >= self._reset_max_cursor:
B
barriery 已提交
791
                    _LOGGER.info(self._log("Reset cursor in Channel"))
792 793 794 795 796 797 798
                    self._base_cursor -= self._reset_max_cursor
                    for name in self._consumer_cursors:
                        self._consumer_cursors[name] -= self._reset_max_cursor
                    self._cursor_count = {
                        cursor - self._reset_max_cursor: count
                        for cursor, count in self._cursor_count.items()
                    }
B
barrierye 已提交
799 800
            else:
                resp = copy.deepcopy(self._output_buf[data_idx])
B
barrierye 已提交
801 802 803 804 805 806

            self._consumer_cursors[op_name] += 1
            new_consumer_cursor = self._consumer_cursors[op_name]
            if self._cursor_count.get(new_consumer_cursor) is None:
                self._cursor_count[new_consumer_cursor] = 0
            self._cursor_count[new_consumer_cursor] += 1
807 808 809

            self._cv.notify_all()

B
barriery 已提交
810
        _LOGGER.debug(
T
TeslaZhao 已提交
811 812 813
            self._log(
                "(data_id={} log_id={}) Op({}) Got data from output_buffer".
                format(resp.values()[0].id, resp.values()[0].log_id, op_name)))
B
barrierye 已提交
814
        return resp
815 816

    def stop(self):
817
        _LOGGER.info(self._log("stop."))
818
        self._stop = True
B
barrierye 已提交
819 820 821
        with self._cv:
            self._cv.notify_all()

B
barriery 已提交
822

B
barriery 已提交
823 824 825
class ChannelTimeoutError(RuntimeError):
    def __init__(self):
        pass
B
barrierye 已提交
826

B
barriery 已提交
827

B
barrierye 已提交
828 829 830
class ChannelStopError(RuntimeError):
    def __init__(self):
        pass