ding.utils.lock_helper

ding.utils.lock_helper

LockContextType

Bases: Enum

Overview

Enum to express the type of the lock.

LockContext

Bases: object

Overview

Generate a LockContext in order to make sure the thread safety.

Interfaces

__init__, __enter__, __exit__.

Example

with LockContext() as lock: print("Do something here.")

__init__(lock_type=LockContextType.THREAD_LOCK)

Overview

Init the lock according to the given type.

Parameters:

Name	Type	Description	Default
- lock_type (		obj:LockContextType): The type of lock to be used. Defaults to LockContextType.THREAD_LOCK.	required

acquire()

Overview

Acquires the lock.

release()

Overview

Releases the lock.

__enter__()

Overview

Enters the context and acquires the lock.

__exit__(*args, **kwargs)

Overview

Exits the context and releases the lock.

Arguments: - args (:obj:Tuple): The arguments passed to the __exit__ function. - kwargs (:obj:Dict): The keyword arguments passed to the __exit__ function.

FcntlContext

Overview

A context manager that acquires an exclusive lock on a file using fcntl. This is useful for preventing multiple processes from running the same code.

Interfaces

__init__, __enter__, __exit__.

Example

lock_path = "/path/to/lock/file" with FcntlContext(lock_path) as lock: # Perform operations while the lock is held

__init__(lock_path)

Overview

Initialize the LockHelper object.

Parameters:

Name	Type	Description	Default
- lock_path (		obj:str): The path to the lock file.	required

__enter__()

Overview

Acquires the lock and opens the lock file in write mode. If the lock file does not exist, it is created.

__exit__(*args, **kwargs)

Overview

Closes the file and releases any resources used by the lock_helper object.

Arguments: - args (:obj:Tuple): The arguments passed to the __exit__ function. - kwargs (:obj:Dict): The keyword arguments passed to the __exit__ function.

get_rw_file_lock(name, op)

Overview

Get generated file lock with name and operator

Arguments: - name (:obj:str): Lock's name. - op (:obj:str): Assigned operator, i.e. read or write. Returns: - (:obj:RWLockFairD): Generated rwlock

get_file_lock(name, op)

Overview

Acquires a file lock for the specified file.

Arguments: - name (:obj:str): The name of the file. - op (:obj:str): The operation to perform on the file lock.

Full Source Code

../ding/utils/lock_helper.py

1import os 2import multiprocessing 3import threading 4import platform 5from enum import Enum, unique 6 7from pathlib import Path 8if platform.system().lower() != 'windows': 9 import fcntl 10else: 11 fcntl = None 12 13 14@unique 15class LockContextType(Enum): 16 """ 17 Overview: 18 Enum to express the type of the lock. 19 """ 20 THREAD_LOCK = 1 21 PROCESS_LOCK = 2 22 23 24_LOCK_TYPE_MAPPING = { 25 LockContextType.THREAD_LOCK: threading.Lock, 26 LockContextType.PROCESS_LOCK: multiprocessing.Lock, 27} 28 29 30class LockContext(object): 31 """ 32 Overview: 33 Generate a LockContext in order to make sure the thread safety. 34 35 Interfaces: 36 ``__init__``, ``__enter__``, ``__exit__``. 37 38 Example: 39 >>> with LockContext() as lock: 40 >>> print("Do something here.") 41 """ 42 43 def __init__(self, lock_type: LockContextType = LockContextType.THREAD_LOCK): 44 """ 45 Overview: 46 Init the lock according to the given type. 47 48 Arguments: 49 - lock_type (:obj:`LockContextType`): The type of lock to be used. Defaults to LockContextType.THREAD_LOCK. 50 """ 51 self.lock = _LOCK_TYPE_MAPPING[lock_type]() 52 53 def acquire(self): 54 """ 55 Overview: 56 Acquires the lock. 57 """ 58 self.lock.acquire() 59 60 def release(self): 61 """ 62 Overview: 63 Releases the lock. 64 """ 65 self.lock.release() 66 67 def __enter__(self): 68 """ 69 Overview: 70 Enters the context and acquires the lock. 71 """ 72 self.lock.acquire() 73 74 def __exit__(self, *args, **kwargs): 75 """ 76 Overview: 77 Exits the context and releases the lock. 78 Arguments: 79 - args (:obj:`Tuple`): The arguments passed to the ``__exit__`` function. 80 - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__exit__`` function. 81 """ 82 self.lock.release() 83 84 85rw_lock_mapping = {} 86 87 88def get_rw_file_lock(name: str, op: str): 89 """ 90 Overview: 91 Get generated file lock with name and operator 92 Arguments: 93 - name (:obj:`str`): Lock's name. 94 - op (:obj:`str`): Assigned operator, i.e. ``read`` or ``write``. 95 Returns: 96 - (:obj:`RWLockFairD`): Generated rwlock 97 """ 98 assert op in ['read', 'write'] 99 try: 100 from readerwriterlock import rwlock 101 except ImportError: 102 import sys 103 from ditk import logging 104 logging.warning("Please install readerwriterlock first, such as `pip3 install readerwriterlock`.") 105 sys.exit(1) 106 if name not in rw_lock_mapping: 107 rw_lock_mapping[name] = rwlock.RWLockFairD() 108 lock = rw_lock_mapping[name] 109 if op == 'read': 110 return lock.gen_rlock() 111 elif op == 'write': 112 return lock.gen_wlock() 113 114 115class FcntlContext: 116 """ 117 Overview: 118 A context manager that acquires an exclusive lock on a file using fcntl. \ 119 This is useful for preventing multiple processes from running the same code. 120 121 Interfaces: 122 ``__init__``, ``__enter__``, ``__exit__``. 123 124 Example: 125 >>> lock_path = "/path/to/lock/file" 126 >>> with FcntlContext(lock_path) as lock: 127 >>> # Perform operations while the lock is held 128 129 """ 130 131 def __init__(self, lock_path: str) -> None: 132 """ 133 Overview: 134 Initialize the LockHelper object. 135 136 Arguments: 137 - lock_path (:obj:`str`): The path to the lock file. 138 """ 139 self.lock_path = lock_path 140 self.f = None 141 142 def __enter__(self) -> None: 143 """ 144 Overview: 145 Acquires the lock and opens the lock file in write mode. \ 146 If the lock file does not exist, it is created. 147 """ 148 assert self.f is None, self.lock_path 149 self.f = open(self.lock_path, 'w') 150 fcntl.flock(self.f.fileno(), fcntl.LOCK_EX) 151 152 def __exit__(self, *args, **kwargs) -> None: 153 """ 154 Overview: 155 Closes the file and releases any resources used by the lock_helper object. 156 Arguments: 157 - args (:obj:`Tuple`): The arguments passed to the ``__exit__`` function. 158 - kwargs (:obj:`Dict`): The keyword arguments passed to the ``__exit__`` function. 159 """ 160 self.f.close() 161 self.f = None 162 163 164def get_file_lock(name: str, op: str) -> FcntlContext: 165 """ 166 Overview: 167 Acquires a file lock for the specified file. \ 168 169 Arguments: 170 - name (:obj:`str`): The name of the file. 171 - op (:obj:`str`): The operation to perform on the file lock. 172 """ 173 if fcntl is None: 174 return get_rw_file_lock(name, op) 175 else: 176 lock_name = name + '.lock' 177 if not os.path.isfile(lock_name): 178 try: 179 Path(lock_name).touch() 180 except Exception as e: 181 pass 182 return FcntlContext(lock_name)