extern/easydiffusion
1
0
Fork 1
mirror of https://github.com/easydiffusion/easydiffusion.git synced 2025-04-24 03:18:29 +02:00
Code Issues Packages Projects Releases Wiki Activity

Added timeout to critical locking tasks with matching exception

Browse Source
This commit is contained in:
Marc-Andre Ferland 2022-10-15 05:31:17 -04:00
parent 7625e591fe
commit cbdf03450d
1 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
ui/sd_internal/task_manager.py
View File

@ -76,7 +76,7 @@ class TaskCache():
def _is_expired(self, timestamp: int) -> bool: def _is_expired(self, timestamp: int) -> bool:
return int(time.time()) >= timestamp return int(time.time()) >= timestamp
def clean(self) -> None: def clean(self) -> None:
self._lock.acquire() if not self._lock.acquire(blocking=True, timeout=10): raise Exception('TaskCache.clean failed to acquire lock within timeout.')
try: try:
# Create a list of expired keys to delete # Create a list of expired keys to delete
to_delete = [] to_delete = []
@ -91,11 +91,11 @@ class TaskCache():
finally: finally:
self._lock.release() self._lock.release()
def clear(self) -> None: def clear(self) -> None:
self._lock.acquire() if not self._lock.acquire(blocking=True, timeout=10): raise Exception('TaskCache.clear failed to acquire lock within timeout.')
try: self._base.clear() try: self._base.clear()
finally: self._lock.release() finally: self._lock.release()
def delete(self, key: Hashable) -> bool: def delete(self, key: Hashable) -> bool:
self._lock.acquire() if not self._lock.acquire(blocking=True, timeout=10): raise Exception('TaskCache.delete failed to acquire lock within timeout.')
try: try:
if key not in self._base: if key not in self._base:
return False return False
@ -104,7 +104,7 @@ class TaskCache():
finally: finally:
self._lock.release() self._lock.release()
def keep(self, key: Hashable, ttl: int) -> bool: def keep(self, key: Hashable, ttl: int) -> bool:
self._lock.acquire() if not self._lock.acquire(blocking=True, timeout=10): raise Exception('TaskCache.keep failed to acquire lock within timeout.')
try: try:
if key in self._base: if key in self._base:
_, value = self._base.get(key) _, value = self._base.get(key)
@ -114,7 +114,7 @@ class TaskCache():
finally: finally:
self._lock.release() self._lock.release()
def put(self, key: Hashable, value: Any, ttl: int) -> bool: def put(self, key: Hashable, value: Any, ttl: int) -> bool:
self._lock.acquire() if not self._lock.acquire(blocking=True, timeout=10): raise Exception('TaskCache.put failed to acquire lock within timeout.')
try: try:
self._base[key] = ( self._base[key] = (
self._get_ttl_time(ttl), value self._get_ttl_time(ttl), value
@ -128,7 +128,7 @@ class TaskCache():
finally: finally:
self._lock.release() self._lock.release()
def tryGet(self, key: Hashable) -> Any: def tryGet(self, key: Hashable) -> Any:
self._lock.acquire() if not self._lock.acquire(blocking=True, timeout=10): raise Exception('TaskCache.tryGet failed to acquire lock within timeout.')
try: try:
ttl, value = self._base.get(key, (None, None)) ttl, value = self._base.get(key, (None, None))
if ttl is not None and self._is_expired(ttl): if ttl is not None and self._is_expired(ttl):
@ -293,6 +293,6 @@ def render(req : ImageRequest):
new_task = RenderTask(r) new_task = RenderTask(r)
if task_cache.put(r.session_id, new_task, TASK_TTL): if task_cache.put(r.session_id, new_task, TASK_TTL):
tasks_queue.put(new_task) tasks_queue.put(new_task, block=True, timeout=30)
return new_task return new_task
raise RuntimeError('Failed to add task to cache.') raise RuntimeError('Failed to add task to cache.')