channel.py 35.8 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


class ChannelDataType(enum.Enum):
60 61 62
    """
    Channel data type
    """
D
dongdaxiang 已提交
63 64 65 66 67
    DICT = 0
    CHANNEL_NPDATA = 1
    ERROR = 2


68 69 70 71
class ChannelData(object):
    def __init__(self,
                 datatype=None,
                 npdata=None,
B
barrierye 已提交
72
                 dictdata=None,
73
                 data_id=None,
T
TeslaZhao 已提交
74 75
                 log_id=None,
                 error_code=None,
B
barrierye 已提交
76
                 error_info=None,
T
TeslaZhao 已提交
77 78
                 prod_error_code=None,
                 prod_error_info=None,
B
barrierye 已提交
79
                 client_need_profile=False):
80 81 82
        '''
        There are several ways to use it:
        
T
TeslaZhao 已提交
83 84 85
        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)
86 87 88 89

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

125 126 127 128 129 130 131 132 133 134 135
    def get_size(self):
        size = 0
        dict_data = None
        if isinstance(self.dictdata, dict):
            for k in self.dictdata:
                size += sys.getsizeof(self.dictdata[k]) + sys.getsizeof(k)
        if isinstance(self.npdata, dict):
            for k in self.npdata:
                size += sys.getsizeof(self.npdata[k]) + sys.getsizeof(k)
        return size

B
barrierye 已提交
136
    def add_profile(self, profile_set):
B
barrierye 已提交
137 138
        if self.client_need_profile is False:
            self.client_need_profile = True
B
barrierye 已提交
139
        self.profile_data_set |= profile_set
140

B
barrierye 已提交
141 142
    @staticmethod
    def check_dictdata(dictdata):
T
TeslaZhao 已提交
143
        error_code = ChannelDataErrcode.OK.value
B
barrierye 已提交
144
        error_info = None
B
barrierye 已提交
145 146 147 148
        if isinstance(dictdata, list):
            # batch data
            for sample in dictdata:
                if not isinstance(sample, dict):
T
TeslaZhao 已提交
149
                    error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
150 151
                    error_info = "Failed to check data: the type of " \
                            "data must be dict, but get {}.".format(type(sample))
B
barrierye 已提交
152 153 154
                    break
        elif not isinstance(dictdata, dict):
            # batch size = 1
T
TeslaZhao 已提交
155
            error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
156 157
            error_info = "Failed to check data: the type of data must " \
                    "be dict, but get {}.".format(type(dictdata))
T
TeslaZhao 已提交
158
        return error_code, error_info
B
barrierye 已提交
159

B
bug fix  
barriery 已提交
160 161
    @staticmethod
    def check_batch_npdata(batch):
T
TeslaZhao 已提交
162
        error_code = ChannelDataErrcode.OK.value
B
bug fix  
barriery 已提交
163 164
        error_info = None
        for npdata in batch:
T
TeslaZhao 已提交
165 166
            error_code, error_info = ChannelData.check_npdata(npdata)
            if error_code != ChannelDataErrcode.OK.value:
B
bug fix  
barriery 已提交
167
                break
T
TeslaZhao 已提交
168
        return error_code, error_info
B
bug fix  
barriery 已提交
169

B
barrierye 已提交
170 171
    @staticmethod
    def check_npdata(npdata):
T
TeslaZhao 已提交
172
        error_code = ChannelDataErrcode.OK.value
173
        error_info = None
W
wangjiawei04 已提交
174 175 176 177
        if isinstance(npdata, list):
            # batch data
            for sample in npdata:
                if not isinstance(sample, dict):
T
TeslaZhao 已提交
178
                    error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
179 180 181
                    error_info = "Failed to check data: the " \
                            "value of data must be dict, but get {}.".format(
                                    type(sample))
W
wangjiawei04 已提交
182 183 184
                    break
                for _, value in sample.items():
                    if not isinstance(value, np.ndarray):
T
TeslaZhao 已提交
185
                        error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
186 187 188
                        error_info = "Failed to check data: the" \
                                " value of data must be np.ndarray, but get {}.".format(
                                        type(value))
T
TeslaZhao 已提交
189
                        return error_code, error_info
W
wangjiawei04 已提交
190 191 192 193
        elif isinstance(npdata, dict):
            # batch_size = 1
            for _, value in npdata.items():
                if not isinstance(value, np.ndarray):
T
TeslaZhao 已提交
194
                    error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
195 196 197
                    error_info = "Failed to check data: the value " \
                            "of data must be np.ndarray, but get {}.".format(
                                    type(value))
W
wangjiawei04 已提交
198 199
                    break
        else:
T
TeslaZhao 已提交
200
            error_code = ChannelDataErrcode.TYPE_ERROR.value
B
barriery 已提交
201 202
            error_info = "Failed to check data: the value of data " \
                    "must be dict, but get {}.".format(type(npdata))
T
TeslaZhao 已提交
203
        return error_code, error_info
204 205 206

    def parse(self):
        feed = None
B
barrierye 已提交
207 208
        if self.datatype == ChannelDataType.CHANNEL_NPDATA.value:
            # return narray
209
            feed = self.npdata
B
barrierye 已提交
210 211 212
        elif self.datatype == ChannelDataType.DICT.value:
            # return dict
            feed = self.dictdata
213
        else:
B
barriery 已提交
214 215
            _LOGGER.critical("Failed to parse channeldata: error " \
                    "type({}) in datatype.".format(self.datatype))
216
            os._exit(-1)
217 218
        return feed

219 220 221 222 223 224 225 226
    def __cmp__(self, other):
        if self.id < other.id:
            return -1
        elif self.id == other.id:
            return 0
        else:
            return 1

227 228
    def get_all_data(self):
        return "type[{}], error_code[{}], data_id[{}], log_id[{}], dict_size[{}]".format(
T
TeslaZhao 已提交
229
            ChannelDataType(self.datatype).name, self.error_code, self.id,
230
            self.log_id, self.get_size())
231 232


B
barrierye 已提交
233
class ProcessChannel(object):
234 235 236 237 238 239 240 241 242
    """ 
    (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 已提交
243
    3. Function front support timeout param to make auto-batching.
244 245 246 247 248

    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 已提交
249 250 251 252 253 254 255 256 257 258 259

    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.
260 261
    """

B
barriery 已提交
262
    def __init__(self, manager, name=None, maxsize=0):
B
barrierye 已提交
263 264 265 266 267 268
        # 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
269
        self._que = manager.PriorityQueue(maxsize=maxsize)
270 271
        self._maxsize = maxsize
        self.name = name
272
        self._stop = manager.Value('i', 0)
273 274 275 276

        self._cv = multiprocessing.Condition()

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

280
        self._reset_max_cursor = 1000000000000000000
B
barrierye 已提交
281 282 283 284
        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()
285

B
barriery 已提交
286 287 288
    def get_maxsize(self):
        return self._maxsize

B
barriery 已提交
289 290 291
    def size(self):
        return self._que.qsize()

292 293 294 295
    def get_producers(self):
        return self._producers

    def get_consumers(self):
B
barrierye 已提交
296
        return self._consumer_cursors.keys()
297 298 299 300 301 302 303

    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:
304
            _LOGGER.critical(
B
barriery 已提交
305 306
                self._log("Failed to add producer: producer({})" \
                        " is already in channel".format(op_name)))
307
            os._exit(-1)
308
        self._producers.append(op_name)
B
barriery 已提交
309
        _LOGGER.debug(self._log("Succ add a producer: {}".format(op_name)))
310 311 312

    def add_consumer(self, op_name):
        """ not thread safe, and can only be called during initialization. """
B
barrierye 已提交
313
        if op_name in self._consumer_cursors:
314
            _LOGGER.critical(
B
barriery 已提交
315 316
                    self._log("Failed to add consumer: consumer({})" \
                            " is already in channel".format(op_name)))
317
            os._exit(-1)
B
barrierye 已提交
318
        self._consumer_cursors[op_name] = 0
319

B
barrierye 已提交
320 321 322
        if self._cursor_count.get(0) is None:
            self._cursor_count[0] = 0
        self._cursor_count[0] += 1
B
barriery 已提交
323
        _LOGGER.debug(self._log("Succ add a consumer: {}".format(op_name)))
324 325

    def push(self, channeldata, op_name=None):
B
barrierye 已提交
326
        _LOGGER.debug(
327 328 329 330
            self._log(
                "(data_id={} log_id={}) Op({}) Enter channel::push producers:{}".
                format(channeldata.id, channeldata.log_id, op_name,
                       len(self._producers))))
331
        if len(self._producers) == 0:
332
            _LOGGER.critical(
333
                self._log(
T
TeslaZhao 已提交
334
                    "(data_id={} log_id={}) Op({}) Failed to push data: expected number"
B
barriery 已提交
335
                    " of producers to be greater than 0, but the it is 0.".
T
TeslaZhao 已提交
336
                    format(channeldata.id, channeldata.log_id, op_name)))
337
            os._exit(-1)
338
        elif len(self._producers) == 1:
339
            start_time = _time()
340
            with self._cv:
341 342
                enter_cv_time = _time()
                push_que_time = enter_cv_time
343
                while self._stop.value == 0:
344
                    try:
T
TeslaZhao 已提交
345 346 347 348
                        self._que.put((channeldata.id, {
                            op_name: channeldata
                        }),
                                      timeout=0)
349
                        push_que_time = _time()
350 351 352
                        break
                    except Queue.Full:
                        self._cv.wait()
353
                if self._stop.value == 1:
B
barrierye 已提交
354
                    raise ChannelStopError()
355
                self._cv.notify_all()
356 357 358 359 360 361 362
                notify_all_time = _time()
                _LOGGER.debug(
                    "(data_id={}) Op({}) channel push cost! enter_cv:{} ms, push_que:{} ms, notify:{} ms, data_size:{}".
                    format(channeldata.id, op_name, (enter_cv_time - start_time)
                           * 1000, (push_que_time - enter_cv_time) * 1000, (
                               notify_all_time - push_que_time) * 1000,
                           channeldata.get_size()))
B
barriery 已提交
363
            _LOGGER.debug(
T
TeslaZhao 已提交
364 365 366
                self._log(
                    "(data_id={} log_id={}) Op({}) Pushed data into internal queue.".
                    format(channeldata.id, channeldata.log_id, op_name)))
367 368
            return True
        elif op_name is None:
369
            _LOGGER.critical(
370
                self._log(
T
TeslaZhao 已提交
371
                    "(data_id={} log_id={}) Op({}) Failed to push data: there are multiple "
B
barriery 已提交
372
                    "producers, so op_name cannot be None.".format(
T
TeslaZhao 已提交
373
                        channeldata.id, channeldata.log_id, op_name)))
374
            os._exit(-1)
375 376 377

        producer_num = len(self._producers)
        data_id = channeldata.id
T
TeslaZhao 已提交
378
        log_id = channeldata.log_id
379 380
        put_data = None
        with self._cv:
B
barrierye 已提交
381 382
            if data_id not in self._input_buf:
                self._input_buf[data_id] = {
383 384 385
                    name: None
                    for name in self._producers
                }
B
barrierye 已提交
386
                self._pushed_producer_count[data_id] = 0
387
            # see: https://docs.python.org/3.6/library/multiprocessing.html?highlight=multiprocess#proxy-objects
B
barrierye 已提交
388 389 390 391 392
            # 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 已提交
393
            if self._pushed_producer_count[data_id] + 1 == producer_num:
B
barrierye 已提交
394 395
                put_data = self._input_buf[data_id]
                self._input_buf.pop(data_id)
B
barrierye 已提交
396
                self._pushed_producer_count.pop(data_id)
397
            else:
B
barrierye 已提交
398
                self._pushed_producer_count[data_id] += 1
399 400

            if put_data is None:
B
barrierye 已提交
401
                _LOGGER.debug(
B
barriery 已提交
402
                    self._log(
T
TeslaZhao 已提交
403 404
                        "(data_id={} log_id={}) Op({}) Pushed data into input_buffer.".
                        format(data_id, log_id, op_name)))
405
            else:
406
                while self._stop.value == 0:
407
                    try:
T
TeslaZhao 已提交
408
                        self._que.put((data_id, put_data), timeout=0)
409 410 411
                        break
                    except Queue.Empty:
                        self._cv.wait()
412
                if self._stop.value == 1:
B
barrierye 已提交
413
                    raise ChannelStopError()
414

B
barrierye 已提交
415
                _LOGGER.debug(
B
barriery 已提交
416
                    self._log(
T
TeslaZhao 已提交
417 418
                        "(data_id={} log_id={}) Op({}) Pushed data into internal_queue.".
                        format(data_id, log_id, op_name)))
419 420 421
            self._cv.notify_all()
        return True

B
barriery 已提交
422
    def front(self, op_name=None, timeout=None):
B
barriery 已提交
423
        _LOGGER.debug(
B
barriery 已提交
424 425
            self._log("Op({}) Getting data[?]; timeout(s)={}".format(op_name,
                                                                     timeout)))
B
barriery 已提交
426
        endtime = None
B
bug fix  
barriery 已提交
427 428 429 430 431
        if timeout is not None:
            if timeout <= 0:
                timeout = None
            else:
                endtime = _time() + timeout
B
barriery 已提交
432

B
barrierye 已提交
433
        if len(self._consumer_cursors) == 0:
434
            _LOGGER.critical(
435
                self._log(
B
barriery 已提交
436 437
                    "Op({}) Failed to get data: expected number of consumers to be " \
                            "greater than 0, but the it is 0.".format(op_name)))
438
            os._exit(-1)
B
barrierye 已提交
439
        elif len(self._consumer_cursors) == 1:
440
            resp = None
441 442 443
            time_1 = int(round(_time() * 1000000))
            time_2 = time_1
            time_3 = time_2
444
            with self._cv:
445
                time_2 = int(round(_time() * 1000000))
446
                while self._stop.value == 0 and resp is None:
447
                    try:
T
TeslaZhao 已提交
448
                        resp = self._que.get(timeout=0)[1]
449
                        time_3 = int(round(_time() * 1000000))
450 451
                        break
                    except Queue.Empty:
B
barriery 已提交
452 453 454
                        if timeout is not None:
                            remaining = endtime - _time()
                            if remaining <= 0.0:
B
barriery 已提交
455
                                _LOGGER.debug(
B
barriery 已提交
456 457
                                    self._log("Op({}) Failed to get data: "
                                              "timeout".format(op_name)))
B
barriery 已提交
458 459 460 461
                                raise ChannelTimeoutError()
                            self._cv.wait(remaining)
                        else:
                            self._cv.wait()
462
                if self._stop.value == 1:
B
barrierye 已提交
463
                    raise ChannelStopError()
464 465 466 467 468 469
            key = list(resp.keys())[0]
            data_id = resp[key].id
            _LOGGER.debug(
                "(data_id={}) op({}) front cost enter_cv:{} ms, queue_get:{} ms".
                format(data_id, op_name, (time_2 - time_1) / 1000.0, (
                    time_3 - time_2) / 1000.0))
T
TeslaZhao 已提交
470 471 472 473 474
            if resp is not None:
                list_values = list(resp.values())
                _LOGGER.debug(
                    self._log("(data_id={} log_id={}) Op({}) Got data".format(
                        list_values[0].id, list_values[0].log_id, op_name)))
475 476
            return resp
        elif op_name is None:
477
            _LOGGER.critical(
478
                self._log(
B
barriery 已提交
479 480
                    "Op({}) Failed to get data: there are multiple consumers, "
                    "so op_name cannot be None.".format(op_name)))
481
            os._exit(-1)
482

B
barrierye 已提交
483 484 485 486 487 488 489 490 491 492
        # 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)
493
        with self._cv:
B
barrierye 已提交
494 495
            # 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.
496
            while self._stop.value == 0 and self._consumer_cursors[
B
barrierye 已提交
497
                    op_name] - self._base_cursor.value >= len(self._output_buf):
498
                try:
T
TeslaZhao 已提交
499
                    channeldata = self._que.get(timeout=0)[1]
B
barrierye 已提交
500
                    self._output_buf.append(channeldata)
T
TeslaZhao 已提交
501
                    list_values = list(channeldata.values())
B
barriery 已提交
502
                    _LOGGER.debug(
B
barriery 已提交
503
                        self._log(
T
TeslaZhao 已提交
504
                            "(data_id={} log_id={}) Op({}) Pop ready item into output_buffer".
T
TeslaZhao 已提交
505 506
                            format(list_values[0].id, list_values[0].log_id,
                                   op_name)))
507 508
                    break
                except Queue.Empty:
B
barriery 已提交
509 510 511
                    if timeout is not None:
                        remaining = endtime - _time()
                        if remaining <= 0.0:
B
barriery 已提交
512
                            _LOGGER.debug(
B
barriery 已提交
513 514
                                self._log("Op({}) Failed to get data: timeout".
                                          format(op_name)))
B
barriery 已提交
515 516 517 518
                            raise ChannelTimeoutError()
                        self._cv.wait(remaining)
                    else:
                        self._cv.wait()
519
            if self._stop.value == 1:
B
barrierye 已提交
520
                raise ChannelStopError()
521

522
            time_1 = int(round(_time() * 1000000))
B
barrierye 已提交
523 524 525 526
            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]
527

B
barrierye 已提交
528 529 530 531 532 533 534 535
            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
536 537
                # to avoid cursor overflow
                if self._base_cursor.value >= self._reset_max_cursor:
B
barriery 已提交
538
                    _LOGGER.info(self._log("Reset cursor in Channel"))
539 540 541 542 543 544 545 546 547 548
                    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 已提交
549 550 551 552 553 554

            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
555

556
            self._cv.notify_all()
557 558
            time_2 = int(round(_time() * 1000000))
            #_LOGGER.warning("self._cv logic cost:{}".format(time2 - time1))
559

T
TeslaZhao 已提交
560 561 562 563 564 565
        if resp is not None:
            list_values = list(resp.values())
            _LOGGER.debug(
                self._log(
                    "(data_id={} log_id={}) Op({}) Got data from output_buffer".
                    format(list_values[0].id, list_values[0].log_id, op_name)))
B
barriery 已提交
566
        return resp
567 568

    def stop(self):
569
        _LOGGER.info(self._log("stop."))
570
        self._stop.value = 1
B
barrierye 已提交
571 572
        with self._cv:
            self._cv.notify_all()
573 574


575
class ThreadChannel(Queue.PriorityQueue):
576 577 578 579 580 581 582 583 584
    """ 
    (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 已提交
585
    3. Function front support timeout param to make auto-batching.
586 587 588 589 590

    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 已提交
591 592 593 594 595 596 597 598 599 600 601

    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.
602 603
    """

B
barriery 已提交
604
    def __init__(self, name=None, maxsize=-1):
605 606 607 608 609 610 611 612
        Queue.Queue.__init__(self, maxsize=maxsize)
        self._maxsize = maxsize
        self.name = name
        self._stop = False

        self._cv = threading.Condition()

        self._producers = []
B
barrierye 已提交
613
        self._pushed_producer_count = {}  # {data_id: count}
B
barrierye 已提交
614
        self._input_buf = {}  # {data_id: {op_name: data}}
615

616
        self._reset_max_cursor = 1000000000000000000
B
barrierye 已提交
617 618 619 620
        self._consumer_cursors = {}  # {op_name: idx}
        self._cursor_count = {}  # {cursor: count}
        self._base_cursor = 0
        self._output_buf = []
621

B
barriery 已提交
622 623 624
    def get_maxsize(self):
        return self._maxsize

B
barriery 已提交
625 626 627
    def size(self):
        return self.qsize()

628 629 630 631
    def get_producers(self):
        return self._producers

    def get_consumers(self):
B
barrierye 已提交
632
        return self._consumer_cursors.keys()
633 634 635 636 637 638 639

    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:
640
            _LOGGER.critical(
B
barriery 已提交
641 642
                self._log("Failed to add producer: producer({}) is "
                          "already in channel".format(op_name)))
643
            os._exit(-1)
644
        self._producers.append(op_name)
B
barriery 已提交
645
        _LOGGER.debug(self._log("Succ add a producer: {}".format(op_name)))
646 647 648

    def add_consumer(self, op_name):
        """ not thread safe, and can only be called during initialization. """
B
barrierye 已提交
649
        if op_name in self._consumer_cursors:
650
            _LOGGER.critical(
B
barriery 已提交
651 652
                self._log("Failed to add consumer: consumer({}) is "
                          "already in channel".format(op_name)))
653
            os._exit(-1)
B
barrierye 已提交
654
        self._consumer_cursors[op_name] = 0
655

B
barrierye 已提交
656 657 658
        if self._cursor_count.get(0) is None:
            self._cursor_count[0] = 0
        self._cursor_count[0] += 1
B
barriery 已提交
659
        _LOGGER.debug(self._log("Succ add a consumer: {}".format(op_name)))
660 661

    def push(self, channeldata, op_name=None):
B
barrierye 已提交
662
        _LOGGER.debug(
T
TeslaZhao 已提交
663 664
            self._log("(data_id={} log_id={}) Op({}) Pushing data".format(
                channeldata.id, channeldata.log_id, op_name)))
665
        if len(self._producers) == 0:
666
            _LOGGER.critical(
667
                self._log(
T
TeslaZhao 已提交
668
                    "(data_id={} log_id={}) Op({}) Failed to push data: expected number of "
B
barriery 已提交
669
                    "producers to be greater than 0, but the it is 0.".format(
T
TeslaZhao 已提交
670
                        channeldata.id, channeldata.log_id, op_name)))
671
            os._exit(-1)
672 673 674 675
        elif len(self._producers) == 1:
            with self._cv:
                while self._stop is False:
                    try:
T
TeslaZhao 已提交
676 677 678 679
                        self.put((channeldata.id, {
                            op_name: channeldata
                        }),
                                 timeout=0)
680 681 682
                        break
                    except Queue.Full:
                        self._cv.wait()
B
barrierye 已提交
683 684
                if self._stop:
                    raise ChannelStopError()
685
                self._cv.notify_all()
B
barriery 已提交
686
            _LOGGER.debug(
T
TeslaZhao 已提交
687 688 689
                self._log(
                    "(data_id={} log_id={}) Op({}) Pushed data into internal_queue.".
                    format(channeldata.id, channeldata.log_id, op_name)))
690 691
            return True
        elif op_name is None:
692
            _LOGGER.critical(
693
                self._log(
T
TeslaZhao 已提交
694
                    "(data_id={} log_id={}) Op({}) Failed to push data: there are multiple"
B
barriery 已提交
695
                    " producers, so op_name cannot be None.".format(
T
TeslaZhao 已提交
696
                        channeldata.id, channeldata.log_id, op_name)))
697
            os._exit(-1)
698 699 700

        producer_num = len(self._producers)
        data_id = channeldata.id
T
TeslaZhao 已提交
701
        log_id = channeldata.log_id
702 703
        put_data = None
        with self._cv:
B
barrierye 已提交
704 705
            if data_id not in self._input_buf:
                self._input_buf[data_id] = {
706 707 708
                    name: None
                    for name in self._producers
                }
B
barrierye 已提交
709
                self._pushed_producer_count[data_id] = 0
B
barrierye 已提交
710
            self._input_buf[data_id][op_name] = channeldata
B
barrierye 已提交
711
            if self._pushed_producer_count[data_id] + 1 == producer_num:
B
barrierye 已提交
712 713
                put_data = self._input_buf[data_id]
                self._input_buf.pop(data_id)
B
barrierye 已提交
714
                self._pushed_producer_count.pop(data_id)
715
            else:
B
barrierye 已提交
716
                self._pushed_producer_count[data_id] += 1
717 718

            if put_data is None:
B
barrierye 已提交
719
                _LOGGER.debug(
B
barriery 已提交
720
                    self._log(
T
TeslaZhao 已提交
721 722
                        "(data_id={} log_id={}) Op({}) Pushed data into input_buffer.".
                        format(data_id, log_id, op_name)))
723 724 725
            else:
                while self._stop is False:
                    try:
T
TeslaZhao 已提交
726
                        self.put((data_id, put_data), timeout=0)
727 728 729
                        break
                    except Queue.Empty:
                        self._cv.wait()
B
barrierye 已提交
730 731
                if self._stop:
                    raise ChannelStopError()
732

B
barrierye 已提交
733
                _LOGGER.debug(
B
barriery 已提交
734
                    self._log(
T
TeslaZhao 已提交
735 736
                        "(data_id={} log_id={}) Op({}) Pushed data into internal_queue.".
                        format(data_id, log_id, op_name)))
737 738 739
            self._cv.notify_all()
        return True

B
barriery 已提交
740
    def front(self, op_name=None, timeout=None):
B
barriery 已提交
741
        _LOGGER.debug(
B
barriery 已提交
742 743
            self._log("Op({}) Getting data[?]; timeout(s)={}".format(op_name,
                                                                     timeout)))
B
barriery 已提交
744
        endtime = None
B
bug fix  
barriery 已提交
745 746 747 748 749
        if timeout is not None:
            if timeout <= 0:
                timeout = None
            else:
                endtime = _time() + timeout
B
barriery 已提交
750

B
barrierye 已提交
751
        if len(self._consumer_cursors) == 0:
752
            _LOGGER.critical(
753
                self._log(
B
barriery 已提交
754 755
                    "Op({}) Failed to get data: expected number of consumers to be "
                    "greater than 0, but the it is 0.".format(op_name)))
756
            os._exit(-1)
B
barrierye 已提交
757
        elif len(self._consumer_cursors) == 1:
758 759 760 761
            resp = None
            with self._cv:
                while self._stop is False and resp is None:
                    try:
T
TeslaZhao 已提交
762
                        resp = self.get(timeout=0)[1]
763 764
                        break
                    except Queue.Empty:
B
barriery 已提交
765 766 767
                        if timeout is not None:
                            remaining = endtime - _time()
                            if remaining <= 0.0:
B
barriery 已提交
768
                                _LOGGER.debug(
B
barriery 已提交
769 770 771
                                    self._log(
                                        "Op({}) Failed to get data: timeout".
                                        format(op_name)))
B
barriery 已提交
772 773 774 775
                                raise ChannelTimeoutError()
                            self._cv.wait(remaining)
                        else:
                            self._cv.wait()
B
barrierye 已提交
776 777
                if self._stop:
                    raise ChannelStopError()
T
TeslaZhao 已提交
778 779 780 781 782
            if resp is not None:
                list_values = list(resp.values())
                _LOGGER.debug(
                    self._log("(data_id={} log_id={}) Op({}) Got data".format(
                        list_values[0].id, list_values[0].log_id, op_name)))
783 784
            return resp
        elif op_name is None:
785
            _LOGGER.critical(
B
barriery 已提交
786 787 788
                self._log("Op({}) Failed to get data: there are multiple "
                          "consumers, so op_name cannot be None.".format(
                              op_name)))
789
            os._exit(-1)
790

B
barrierye 已提交
791 792 793 794 795 796 797 798 799 800
        # 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)
801
        with self._cv:
B
barrierye 已提交
802 803 804 805
            # 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):
806
                try:
T
TeslaZhao 已提交
807
                    channeldata = self.get(timeout=0)[1]
B
barrierye 已提交
808
                    self._output_buf.append(channeldata)
T
TeslaZhao 已提交
809
                    list_values = list(channeldata.values())
B
barriery 已提交
810
                    _LOGGER.debug(
B
barriery 已提交
811
                        self._log(
T
TeslaZhao 已提交
812
                            "(data_id={} log_id={}) Op({}) Pop ready item into output_buffer".
T
TeslaZhao 已提交
813 814
                            format(list_values[0].id, list_values[0].log_id,
                                   op_name)))
815 816
                    break
                except Queue.Empty:
B
barriery 已提交
817 818 819
                    if timeout is not None:
                        remaining = endtime - _time()
                        if remaining <= 0.0:
B
barriery 已提交
820
                            _LOGGER.debug(
B
barriery 已提交
821 822
                                self._log("Op({}) Failed to get data: timeout".
                                          format(op_name)))
B
barriery 已提交
823 824 825 826
                            raise ChannelTimeoutError()
                        self._cv.wait(remaining)
                    else:
                        self._cv.wait()
B
barrierye 已提交
827 828
            if self._stop:
                raise ChannelStopError()
829

B
barrierye 已提交
830 831 832
            consumer_cursor = self._consumer_cursors[op_name]
            base_cursor = self._base_cursor
            data_idx = consumer_cursor - base_cursor
B
barrierye 已提交
833 834

            resp = None
835

B
barrierye 已提交
836 837 838 839 840 841
            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 已提交
842
                resp = self._output_buf.pop(0)
B
barrierye 已提交
843
                self._base_cursor += 1
844 845
                # to avoid cursor overflow
                if self._base_cursor >= self._reset_max_cursor:
B
barriery 已提交
846
                    _LOGGER.info(self._log("Reset cursor in Channel"))
847 848 849 850 851 852 853
                    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 已提交
854 855
            else:
                resp = copy.deepcopy(self._output_buf[data_idx])
B
barrierye 已提交
856 857 858 859 860 861

            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
862 863 864

            self._cv.notify_all()

T
TeslaZhao 已提交
865 866 867 868 869 870
        if resp is not None:
            list_values = list(resp.values())
            _LOGGER.debug(
                self._log(
                    "(data_id={} log_id={}) Op({}) Got data from output_buffer".
                    format(list_values[0].id, list_values[0].log_id, op_name)))
B
barrierye 已提交
871
        return resp
872 873

    def stop(self):
874
        _LOGGER.info(self._log("stop."))
875
        self._stop = True
B
barrierye 已提交
876 877 878
        with self._cv:
            self._cv.notify_all()

B
barriery 已提交
879

B
barriery 已提交
880 881 882
class ChannelTimeoutError(RuntimeError):
    def __init__(self):
        pass
B
barrierye 已提交
883

B
barriery 已提交
884

B
barrierye 已提交
885 886 887
class ChannelStopError(RuntimeError):
    def __init__(self):
        pass