Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,204 Bytes
d1ed09d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 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 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
from __future__ import annotations
import asyncio
import logging
import uuid
from collections import defaultdict, deque
from dask.utils import parse_timedelta
from distributed.client import Client
from distributed.utils import TimeoutError, log_errors
from distributed.worker import get_worker
logger = logging.getLogger(__name__)
class LockExtension:
"""An extension for the scheduler to manage Locks
This adds the following routes to the scheduler
* lock_acquire
* lock_release
"""
def __init__(self, scheduler):
self.scheduler = scheduler
self.events = defaultdict(deque)
self.ids = dict()
self.scheduler.handlers.update(
{"lock_acquire": self.acquire, "lock_release": self.release}
)
@log_errors
async def acquire(self, name=None, id=None, timeout=None):
if isinstance(name, list):
name = tuple(name)
if name not in self.ids:
result = True
else:
while name in self.ids:
event = asyncio.Event()
self.events[name].append(event)
future = event.wait()
if timeout is not None:
future = asyncio.wait_for(future, timeout)
try:
await future
except TimeoutError:
result = False
break
else:
result = True
finally:
event2 = self.events[name].popleft()
assert event is event2
if result:
assert name not in self.ids
self.ids[name] = id
return result
@log_errors
def release(self, name=None, id=None):
if isinstance(name, list):
name = tuple(name)
if self.ids.get(name) != id:
raise ValueError("This lock has not yet been acquired")
del self.ids[name]
if self.events[name]:
self.scheduler.loop.add_callback(self.events[name][0].set)
else:
del self.events[name]
class Lock:
"""Distributed Centralized Lock
Parameters
----------
name: string (optional)
Name of the lock to acquire. Choosing the same name allows two
disconnected processes to coordinate a lock. If not given, a random
name will be generated.
client: Client (optional)
Client to use for communication with the scheduler. If not given, the
default global client will be used.
Examples
--------
>>> lock = Lock('x') # doctest: +SKIP
>>> lock.acquire(timeout=1) # doctest: +SKIP
>>> # do things with protected resource
>>> lock.release() # doctest: +SKIP
"""
def __init__(self, name=None, client=None):
try:
self.client = client or Client.current()
except ValueError:
# Initialise new client
self.client = get_worker().client
self.name = name or "lock-" + uuid.uuid4().hex
self.id = uuid.uuid4().hex
self._locked = False
def acquire(self, blocking=True, timeout=None):
"""Acquire the lock
Parameters
----------
blocking : bool, optional
If false, don't wait on the lock in the scheduler at all.
timeout : string or number or timedelta, optional
Seconds to wait on the lock in the scheduler. This does not
include local coroutine time, network transfer time, etc..
It is forbidden to specify a timeout when blocking is false.
Instead of number of seconds, it is also possible to specify
a timedelta in string format, e.g. "200ms".
Examples
--------
>>> lock = Lock('x') # doctest: +SKIP
>>> lock.acquire(timeout="1s") # doctest: +SKIP
Returns
-------
True or False whether or not it successfully acquired the lock
"""
timeout = parse_timedelta(timeout)
if not blocking:
if timeout is not None:
raise ValueError("can't specify a timeout for a non-blocking call")
timeout = 0
result = self.client.sync(
self.client.scheduler.lock_acquire,
name=self.name,
id=self.id,
timeout=timeout,
)
self._locked = True
return result
def release(self):
"""Release the lock if already acquired"""
if not self.locked():
raise ValueError("Lock is not yet acquired")
result = self.client.sync(
self.client.scheduler.lock_release, name=self.name, id=self.id
)
self._locked = False
return result
def locked(self):
return self._locked
def __enter__(self):
self.acquire()
return self
def __exit__(self, exc_type, exc_value, traceback):
self.release()
async def __aenter__(self):
await self.acquire()
return self
async def __aexit__(self, exc_type, exc_value, traceback):
await self.release()
def __reduce__(self):
return (Lock, (self.name,))
|