parallel_map.py 9.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
# Copyright (c) 2019 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.

# function:
16
#   transform samples in 'source' using 'worker'
17 18 19 20 21

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

22
import os
23 24
import sys
import six
25 26 27 28 29
if six.PY3:
    from queue import Empty
else:
    from Queue import Empty

30 31 32 33 34 35 36
import uuid
import logging
import signal
import threading

logger = logging.getLogger(__name__)

37
main_pid = os.getpid()
K
Kaipeng Deng 已提交
38 39
worker_set = set()

40 41

class EndSignal(object):
42 43
    """ signal used to notify worker to exit
    """
44

45 46
    def __init__(self, id, errno=0, errmsg=''):
        self.id = id
47 48 49 50
        self.errno = errno
        self.errmsg = errmsg


51
class ParallelMap(object):
52
    """
53 54 55
    Transform samples to mapped samples which is similar to 
    'basic.MappedDataset', but multiple workers (threads or processes) 
    will be used
56 57 58 59 60

    Notes:
        this class is not thread-safe
    """

61 62 63 64 65 66 67 68 69 70
    def __init__(self,
                 source,
                 worker,
                 worker_num,
                 bufsize=100,
                 use_process=False,
                 memsize='3G'):
        self._worker_num = worker_num
        self._bufsize = bufsize
        self._use_process = use_process
71
        if self._use_process and sys.platform == "win32":
Y
Yang Zhang 已提交
72 73
            logger.debug("Use multi-thread reader instead of "
                         "multi-process reader on Windows.")
74
            self._use_process = False
75
        if self._use_process and type(memsize) is str:
K
Kaipeng Deng 已提交
76 77 78 79 80
            assert memsize[-1].lower() in ['g', 'm'], \
                "invalid param for memsize[%s], should be " \
                "ended with 'G' or 'g' or 'M' or 'm'" % (memsize)
            power = 3 if memsize[-1].lower() == 'g' else 2
            self._memsize = int(memsize[:-1]) * (1024**power)
81 82
        self._started = False
        self._source = source
83
        self._worker = worker
84 85
        self._exit = False
        self._setup()
86 87 88 89 90 91 92
        self._souce_drained = False

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()
93 94 95

    def _setup(self):
        """setup input/output queues and workers """
96
        use_process = self._use_process
97

98
        bufsize = self._bufsize
99 100 101 102
        if use_process:
            from .shared_queue import SharedQueue as Queue
            from multiprocessing import Process as Worker
            from multiprocessing import Event
103
            memsize = self._memsize
W
walloollaw 已提交
104 105
            self._inq = Queue(bufsize, memsize=memsize)
            self._outq = Queue(bufsize, memsize=memsize)
106 107 108 109 110 111 112
        else:
            if six.PY3:
                from queue import Queue
            else:
                from Queue import Queue
            from threading import Thread as Worker
            from threading import Event
W
walloollaw 已提交
113 114
            self._inq = Queue(bufsize)
            self._outq = Queue(bufsize)
115

116
        consumer_num = self._worker_num
117 118 119 120 121 122 123
        id = str(uuid.uuid4())[-3:]
        self._producer = threading.Thread(
            target=self._produce,
            args=('producer-' + id, self._source, self._inq))
        self._producer.daemon = True

        self._consumers = []
124
        self._consumer_endsig = {}
K
Kaipeng Deng 已提交
125
        global worker_set
126
        for i in range(consumer_num):
127
            consumer_id = 'consumer-' + id + '-' + str(i)
128 129
            p = Worker(
                target=self._consume,
130
                args=(consumer_id, self._inq, self._outq, self._worker))
131 132
            self._consumers.append(p)
            p.daemon = True
133
            setattr(p, 'id', consumer_id)
134 135
            if use_process:
                worker_set.add(p)
136 137 138 139 140 141 142 143

        self._epoch = -1
        self._feeding_ev = Event()
        self._produced = 0  # produced sample in self._produce
        self._consumed = 0  # consumed sample in self.next

    def _produce(self, id, source, inq):
        """Fetch data from source and feed it to 'inq' queue"""
144
        endsig = EndSignal(id)
145 146 147 148 149
        while True:
            self._feeding_ev.wait()
            if self._exit:
                break
            try:
150 151
                s = source.next()
                inq.put(s)
152 153
                self._produced += 1
            except StopIteration:
154
                self._souce_drained = True
155
                self._feeding_ev.clear()
156
                self._feeding_ev.wait()
157
            except Exception as e:
158 159 160 161
                endsig.errno = -1
                endsig.errmsg = "producer[{}] failed with error: {}" \
                    .format(id, str(e))
                inq.put(endsig)
162 163
                break

164
    def _consume(self, id, inq, outq, worker):
165
        """Fetch data from 'inq', process it and put result to 'outq'"""
166
        if self._use_process:
167
            # handle SIGTERM signal to exit to prevent print stack frame
168
            signal.signal(signal.SIGTERM, lambda signum, frame: sys.exit())
169 170

        endsig = EndSignal(id)
171 172 173
        while True:
            sample = inq.get()
            if isinstance(sample, EndSignal):
174 175 176 177
                endsig.errno = sample.errno
                endsig.errmsg = "consumer[{}] exits for reason[{}]" \
                    .format(id, sample.errmsg)
                outq.put(endsig)
178 179 180
                break

            try:
181
                result = worker(sample)
182 183
                outq.put(result)
            except Exception as e:
184 185 186 187
                endsig.errno = -2
                endsig.errmsg = "consumer[{}] failed to map with error:[{}]" \
                    .format(id, str(e))
                outq.put(endsig)
188 189 190 191 192 193 194 195 196 197 198 199 200 201
                break

    def drained(self):
        assert self._epoch >= 0, "first epoch has not started yet"
        return self._source.drained() and self._produced == self._consumed

    def stop(self):
        """ notify to exit
        """
        self._exit = True
        self._feeding_ev.set()
        for _ in range(len(self._consumers)):
            self._inq.put(EndSignal(0, "notify consumers to exit"))

202 203 204 205 206
    def _consumer_healthy(self):
        abnormal_num = 0
        for w in self._consumers:
            if not w.is_alive() and w.id not in self._consumer_endsig:
                abnormal_num += 1
207
                if self._use_process:
208 209 210 211
                    errmsg = "consumer[{}] exit abnormally with exitcode[{}]" \
                                .format(w.pid, w.exitcode)
                else:
                    errmsg = "consumer[{}] exit abnormally".format(w.ident)
212

213 214 215 216 217 218 219 220
                logger.warn(errmsg)

        if abnormal_num > 0:
            logger.warn("{} consumers have exited abnormally!!!" \
                .format(abnormal_num))

        return abnormal_num == 0

221 222 223 224 225 226 227 228 229
    def next(self):
        """ get next transformed sample
        """
        if self._epoch < 0:
            self.reset()

        if self.drained():
            raise StopIteration()

230 231 232 233 234 235 236 237 238
        while not self._exit:
            try:
                sample = self._outq.get(timeout=3)
            except Empty as e:
                if not self._consumer_healthy():
                    raise StopIteration()
                else:
                    continue

239
            if isinstance(sample, EndSignal):
240 241 242
                self._consumer_endsig[sample.id] = sample
                logger.warn("recv endsignal from outq with errmsg[{}]" \
                    .format(sample.errmsg))
243

244
                if len(self._consumer_endsig.keys()) < len(self._consumers):
245 246
                    self._inq.put(sample)
                else:
247 248
                    self._exit = True
                    raise StopIteration("all consumers exited, no more samples")
249 250 251 252
            else:
                self._consumed += 1
                return sample

253 254
        raise StopIteration()

255 256 257
    def reset(self):
        """ reset for a new epoch of samples
        """
258 259
        assert not self._exit, "cannot reset for already stopped dataset"

260 261
        if self._epoch < 0:
            self._epoch = 0
262 263
            for w in self._consumers:
                w.start()
264 265
            self._producer.start()
        else:
266 267 268
            assert self._consumer_healthy(), "cannot start another pass of data" \
                " for some consumers exited abnormally before!!!"

269
            if not self.drained():
270 271
                logger.warn("reset before epoch[{}] finishes".format(
                    self._epoch))
272 273 274 275 276 277
                self._produced = self._produced - self._consumed
            else:
                self._produced = 0

            self._epoch += 1

278
        assert len(self._consumer_endsig.keys()) == 0, "some consumers already exited," \
279 280 281
            + " cannot start another epoch"

        self._source.reset()
282
        self._souce_drained = False
283 284 285 286
        self._consumed = 0
        self._feeding_ev.set()


K
Kaipeng Deng 已提交
287
# FIXME: fix me if you have better impliment
288
# handle terminate reader process, do not print stack frame
289
signal.signal(signal.SIGTERM, lambda signum, frame: sys.exit())
290 291


292 293 294 295
# FIXME(dkp): KeyboardInterrupt should be handled inside ParallelMap
# and do such as: 1. exit workers 2. close queues 3. release shared
# memory, HACK KeyboardInterrupt with global signal.SIGINT handler
# here, should be refined later
K
Kaipeng Deng 已提交
296
def _term_workers(sig_num, frame):
297 298 299 300 301 302 303
    global worker_set, main_pid
    # only do subporcess killing in main process
    if os.getpid() != main_pid:
        return

    logger.info("KeyboardInterrupt: main proc {} exit, kill subprocess {}" \
                .format(os.getpid(), [w.pid for w in worker_set]))
K
Kaipeng Deng 已提交
304
    for w in worker_set:
305 306 307
        if w.pid is not None:
            os.kill(w.pid, signal.SIGINT)
    sys.exit()
308 309


K
Kaipeng Deng 已提交
310
signal.signal(signal.SIGINT, _term_workers)