locks.py 12.4 KB
Newer Older
1
2
3
4
# -*- coding: utf-8 -*-
__author__ = 'Daniel Scheffler'

import time
5
from redis_lock import StrictRedis, Lock, NotAcquired, AlreadyAcquired
6
from redis_semaphore import Semaphore
Daniel Scheffler's avatar
Daniel Scheffler committed
7
from redis.exceptions import ConnectionError as RedisConnectionError
8
import functools
9
import re
10
import random
11
from psutil import virtual_memory
Daniel Scheffler's avatar
Daniel Scheffler committed
12
import timeout_decorator
13
14
15
16

from ..misc.logging import GMS_logger
from ..options.config import GMS_config as CFG

17
18

try:
19
    redis_conn = StrictRedis(host='localhost', db=0)
20
    redis_conn.keys()  # may raise ConnectionError, e.g., if redis server is not installed or not running
Daniel Scheffler's avatar
Daniel Scheffler committed
21
except RedisConnectionError:
22
23
24
    redis_conn = None


25
class MultiSlotLock(Lock):
Daniel Scheffler's avatar
Daniel Scheffler committed
26
    def __init__(self, name, allowed_slots=1, logger=None, disabled=False, **kwargs):
27
        self.conn = redis_conn
28
        self.name = name
29
        self.allowed_slots = allowed_slots or 0
30
31
        self.logger = logger or GMS_logger("RedisLock: '%s'" % self.name)
        self.kwargs = kwargs
32

33
        self.allowed_slot_names = ['%s, slot #%s' % (self.name, i) for i in range(1, self.allowed_slots + 1)]
34
        self.final_name = ''
35
        self._acquired = None
36

Daniel Scheffler's avatar
Daniel Scheffler committed
37
        if not disabled and allowed_slots and redis_conn:
38
39
            logged = False
            while True:
40
                time.sleep(random.uniform(0, 1.5))  # avoids race conditions in case multiple tasks are waiting
41
                name_free_slot = self.get_free_slot_name()
42

43
                if not name_free_slot:
44
                    if not logged:
45
                        self.logger.info("Waiting for free '%s' lock." % self.name)
46
47
48
                        logged = True
                else:
                    break
49

50
51
            self.final_name = 'GMS_%s__' % CFG.ID + name_free_slot
            super().__init__(self.conn, self.final_name, **kwargs)
52
53
54
        else:
            pass

55
56
    @property
    def existing_locks(self):
57
58
59
        names = [i.decode('utf8').split('lock:')[1] for i in self.conn.keys() if i.decode('utf8').startswith('lock:')]

        # split 'GMS_<jobid>' and return
60
        return list(set([re.search('GMS_[0-9]*__(.*)', n, re.I).group(1) for n in names if n.startswith('GMS_')]))
61
62

    def get_free_slot_name(self):
63
        free_slots = [sn for sn in self.allowed_slot_names if sn not in self.existing_locks]
64
65
66
        if free_slots:
            return free_slots[0]

67
    def acquire(self, blocking=True, timeout=None):
68
        if self.allowed_slots and self.conn:
69
70
71
72
73
            if self._acquired:
                raise AlreadyAcquired("Already acquired from this Lock instance.")

            while not self._acquired:
                try:
74
                    # print('Trying to acquire %s.' % self.final_name.split('GMS_%s__' % CFG.ID)[1])
75
                    self._acquired = super(MultiSlotLock, self).acquire(blocking=blocking, timeout=timeout)
76
                    # print("Acquired lock '%s'." % self.final_name.split('GMS_%s__' % CFG.ID)[1])
77
78
79
80
                except AlreadyAcquired:
                    # this happens in case the lock has already been acquired by another instance of MultiSlotLock due
                    # to a race condition (time gap between finding the free slot and the call of self.acquire())
                    # -> in that case: re-initialize to get a new free slot
81
                    self.__init__(self.name, allowed_slots=self.allowed_slots, logger=self.logger,
82
83
84
                                  **self.kwargs)

                if self._acquired is False:  # and not None
85
                    self.__init__(self.name, allowed_slots=self.allowed_slots, logger=self.logger,
86
87
                                  **self.kwargs)

88
89
                # print(self.final_name.split('GMS_%s__' % CFG.ID)[1], self._acquired)

90
91
            if self._acquired:
                self.logger.info("Acquired lock '%s'." % self.final_name.split('GMS_%s__' % CFG.ID)[1])
92
        else:
93
            self._acquired = True
94

95
        return self._acquired
96

97
    def release(self):
98
        if self.allowed_slots and self.conn:
99
100
            super(MultiSlotLock, self).release()
            self.logger.info("Released lock '%s'." % self.final_name.split('GMS_%s__' % CFG.ID)[1])
101
102


103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
class MultiSlotLock(Semaphore):
    def __init__(self, name='MultiSlotLock', allowed_slots=1, logger=None, **kwargs):
        self.disabled = redis_conn is None
        self.name = name
        self.redis_varname = 'GMS_%s__' % CFG.ID + name
        self.allowed_slots = allowed_slots
        self.logger = logger or GMS_logger("RedisLock: '%s'" % name)

        if not self.disabled:
            super().__init__(client=redis_conn, count=allowed_slots, namespace=self.redis_varname, **kwargs)

    def acquire(self, timeout=0, target=None):
        if not self.disabled:
            if self.available_count == 0:
                self.logger.info("Waiting for free lock '%s'." % self.name)

            token = super().acquire(timeout=timeout, target=target)

            self.logger.info("Acquired lock '%s'" % self.name +
                             ('.' if self.allowed_slots == 1 else ', slot #%s.' % (int(token) + 1)))

            return token

    def release(self):
        if not self.disabled:
            token = super().release()
            if token:
                self.logger.info("Released lock '%s'" % self.name +
                                 ('.' if self.allowed_slots == 1 else ', slot #%s.' % (int(token) + 1)))

    def delete(self):
        if not self.disabled:
            redis_conn.delete('%s:EXISTS' % self.redis_varname)
            redis_conn.delete('%s:AVAILABLE' % self.redis_varname)
            redis_conn.delete('%s:GRABBED' % self.redis_varname)


140
class IOLock(MultiSlotLock):
141
142
    def __init__(self, allowed_slots=1, logger=None, **kwargs):
        super(IOLock, self).__init__(name='IOLock', allowed_slots=allowed_slots, logger=logger, **kwargs)
143
144


145
class ProcessLock(MultiSlotLock):
146
147
    def __init__(self, allowed_slots=1, logger=None, **kwargs):
        super(ProcessLock, self).__init__(name='ProcessLock', allowed_slots=allowed_slots, logger=logger, **kwargs)
148
149


Daniel Scheffler's avatar
Daniel Scheffler committed
150
151
class MemoryLock(MultiSlotLock):
    def __init__(self, allowed_slots=1, usage_threshold=80, blocking_threshold=90, logger=None, **kwargs):
152
        """
153

154
155
156
157
158
159
        :param allowed_slots:
        :param usage_th:        Memory usage threshold above which the memory lock is active (percent).
        :param logger:
        :param kwargs:
        """
        self.usage_th = usage_threshold
Daniel Scheffler's avatar
Daniel Scheffler committed
160
        self.blocking_th = blocking_threshold
161
        self.logger = logger or GMS_logger("RedisLock: '%s'" % self.name)
Daniel Scheffler's avatar
Daniel Scheffler committed
162
        self.mem_usage = virtual_memory().percent
163
164
165

        self._acquired = None

Daniel Scheffler's avatar
Daniel Scheffler committed
166
167
        self.check_system_overload()

168
        if self.mem_usage >= usage_threshold:
Daniel Scheffler's avatar
Daniel Scheffler committed
169
170
            self.logger.info('Memory usage is %.1f percent. Acquiring memory lock..' % self.mem_usage)
            self.disabled = False
171
        else:
Daniel Scheffler's avatar
Daniel Scheffler committed
172
173
174
175
176
177
178
179
180
181
182
183
184
185
            self.logger.debug('Memory usage is %.1f percent. Memory lock not needed -> skipped.' % self.mem_usage)
            self.disabled = True

        super(MemoryLock, self).__init__(name='MemoryLock', allowed_slots=allowed_slots, logger=self.logger,
                                         disabled=self.disabled, **kwargs)

    @timeout_decorator.timeout(seconds=15*60, timeout_exception=TimeoutError,
                               exception_message='Memory lock could not be acquired after waiting 15 minutes '
                                                 'due to memory overload.')
    def check_system_overload(self):
        logged = False
        while self.mem_usage >= self.blocking_th:
            if not logged:
                self.logger.info('Memory usage is %.1f percent. Waiting until memory usage is below %s percent.'
Daniel Scheffler's avatar
Bugfix.    
Daniel Scheffler committed
186
                                 % (self.mem_usage, self.blocking_th))
Daniel Scheffler's avatar
Bugfix.    
Daniel Scheffler committed
187
                logged = True
Daniel Scheffler's avatar
Daniel Scheffler committed
188
189
            time.sleep(5)
            self.mem_usage = virtual_memory().percent
190
191
192

    def acquire(self, blocking=True, timeout=None):
        if self.mem_usage > self.usage_th:
193
194
            # self._acquired = super(MemoryLock, self).acquire(blocking=blocking, timeout=timeout)
            self._acquired = super(MemoryLock, self).acquire(timeout=timeout)
195
196
197
198
199
200
        else:
            self._acquired = True

        return self._acquired

    def release(self):
Daniel Scheffler's avatar
Daniel Scheffler committed
201
202
        if self._acquired and not self.disabled:
            super(MemoryLock, self).release()
203
204


205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
class MemoryLock(MultiSlotLock):
    def __init__(self, name='MemoryLock', allowed_slots=1, usage_threshold=80, blocking_threshold=90, logger=None,
                 **kwargs):
        """

        :param allowed_slots:
        :param usage_th:        Memory usage threshold above which the memory lock is active (percent).
        :param logger:
        :param kwargs:
        """
        self.usage_th = usage_threshold
        self.blocking_th = blocking_threshold
        self.logger = logger or GMS_logger("RedisLock: '%s'" % self.name)
        self.mem_usage = virtual_memory().percent

        self._acquired = None

        self.check_system_overload()

        if self.mem_usage >= usage_threshold:
            self.logger.info('Memory usage is %.1f percent. Acquiring memory lock..' % self.mem_usage)
            self.disabled = False
        else:
            self.logger.debug('Memory usage is %.1f percent. Memory lock not needed -> skipped.' % self.mem_usage)
            self.disabled = True

        super(MemoryLock, self).__init__(name='MemoryLock', allowed_slots=allowed_slots, logger=self.logger,
                                         disabled=self.disabled, **kwargs)

    @timeout_decorator.timeout(seconds=15*60, timeout_exception=TimeoutError,
                               exception_message='Memory lock could not be acquired after waiting 15 minutes '
                                                 'due to memory overload.')
    def check_system_overload(self):
        logged = False
        while self.mem_usage >= self.blocking_th:
            if not logged:
                self.logger.info('Memory usage is %.1f percent. Waiting until memory usage is below %s percent.'
                                 % (self.mem_usage, self.blocking_th))
                logged = True
            time.sleep(5)
            self.mem_usage = virtual_memory().percent

    def acquire(self, blocking=True, timeout=None):
        if self.mem_usage > self.usage_th:
            # self._acquired = super(MemoryLock, self).acquire(blocking=blocking, timeout=timeout)
            self._acquired = super(MemoryLock, self).acquire(timeout=timeout)
        else:
            self._acquired = True

        return self._acquired

    def release(self):
        if self._acquired and not self.disabled:
            super(MemoryLock, self).release()



262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
def acquire_process_lock(**processlock_kwargs):
    """Decorator function for ProcessLock.

    :param processlock_kwargs:  Keywourd arguments to be passed to ProcessLock class.
    """
    def decorator(func):

        @functools.wraps(func)  # needed to avoid pickling errors
        def wrapped_func(*args, **kwargs):
            with ProcessLock(**processlock_kwargs):
                result = func(*args, **kwargs)

            return result

        return wrapped_func

    return decorator


def acquire_mem_lock(**memlock_kwargs):
    """Decorator function for MEMLock.

Daniel Scheffler's avatar
Daniel Scheffler committed
284
    :param memlock_kwargs:  Keywourd arguments to be passed to MemoryLock class.
285
    """
286
287
288
289
    def decorator(func):

        @functools.wraps(func)  # needed to avoid pickling errors
        def wrapped_func(*args, **kwargs):
Daniel Scheffler's avatar
Daniel Scheffler committed
290
            with MemoryLock(**memlock_kwargs):
291
292
293
294
295
296
297
298
299
300
301
302
303
                result = func(*args, **kwargs)

            return result

        return wrapped_func

    return decorator


def release_unclosed_locks(logger=None):
    if redis_conn:
        logger = logger or GMS_logger('LockReseter')

304
305
306
        IOLock(allowed_slots=1).delete()
        ProcessLock(allowed_slots=1).delete()

307
308
309
310
311
312
313
314
315
316
317
        locks2release = [i.split(b'lock:')[1].decode('latin') for i in redis_conn.keys()
                         if i.decode('latin').startswith('lock:GMS_%s__' % CFG.ID)]
        if locks2release:
            logger.info("Releasing unclosed locks of job %s." % CFG.ID)

        for lockN in locks2release:
            lock = Lock(redis_conn, lockN)
            try:
                lock.release()
            except NotAcquired:
                lock.reset()
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333


class MemoryReserver(object):
    def __init__(self, reserved_mem=10):
        """

        :param reserved_mem:    Amount of memory to be reserved during the lock is acquired (gigabytes).
        """

    def __enter__(self):
        # while

        return self

    def __exit__(self, exc_type=None, exc_value=None, traceback=None):
        """"""