98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Coordinate quant cache generation across threads and ComfyUI processes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
import time
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
|
|
LOCK_POLL_SECONDS = 0.2
|
|
LOCK_TIMEOUT_SECONDS = 60 * 60
|
|
STALE_LOCK_SECONDS = 2 * 60 * 60
|
|
|
|
|
|
class QuantCacheBuildCoordinator:
|
|
"""Serialize builders for the same derived checkpoint identity."""
|
|
|
|
_thread_locks: dict[str, threading.Lock] = {}
|
|
_thread_locks_guard = threading.Lock()
|
|
|
|
def __init__(self, lock_directory: Path) -> None:
|
|
"""Create a coordinator rooted in the global cache directory."""
|
|
|
|
self._lock_directory = lock_directory
|
|
|
|
@contextmanager
|
|
def acquire(self, stable_key: str) -> Iterator[None]:
|
|
"""Acquire in-process and cross-process ownership for one cache key."""
|
|
|
|
thread_lock = self._thread_lock(stable_key)
|
|
with thread_lock:
|
|
self._lock_directory.mkdir(parents=True, exist_ok=True)
|
|
lock_path = self._lock_directory / f"{stable_key}.lock"
|
|
self._acquire_file_lock(lock_path)
|
|
try:
|
|
yield
|
|
finally:
|
|
try:
|
|
lock_path.unlink(missing_ok=True)
|
|
except OSError:
|
|
pass
|
|
|
|
@classmethod
|
|
def _thread_lock(cls, stable_key: str) -> threading.Lock:
|
|
"""Return the shared in-process lock for a stable cache key."""
|
|
|
|
with cls._thread_locks_guard:
|
|
return cls._thread_locks.setdefault(stable_key, threading.Lock())
|
|
|
|
@staticmethod
|
|
def _acquire_file_lock(lock_path: Path) -> None:
|
|
"""Create an exclusive lock file, recovering only demonstrably stale locks."""
|
|
|
|
started = time.monotonic()
|
|
while True:
|
|
try:
|
|
descriptor = os.open(
|
|
lock_path,
|
|
os.O_CREAT | os.O_EXCL | os.O_WRONLY,
|
|
)
|
|
with os.fdopen(descriptor, "w", encoding="utf-8") as lock_file:
|
|
lock_file.write(f"pid={os.getpid()}\ncreated={time.time()}\n")
|
|
return
|
|
except FileExistsError:
|
|
try:
|
|
age = time.time() - lock_path.stat().st_mtime
|
|
if age > STALE_LOCK_SECONDS and not _lock_owner_is_alive(lock_path):
|
|
lock_path.unlink()
|
|
continue
|
|
except FileNotFoundError:
|
|
continue
|
|
if time.monotonic() - started >= LOCK_TIMEOUT_SECONDS:
|
|
raise TimeoutError(
|
|
"Timed out waiting for another SimpleSyrup process to finish "
|
|
"the same quantized checkpoint."
|
|
) from None
|
|
time.sleep(LOCK_POLL_SECONDS)
|
|
|
|
|
|
def _lock_owner_is_alive(lock_path: Path) -> bool:
|
|
"""Return whether a well-formed lock still belongs to a running process."""
|
|
|
|
try:
|
|
first_line = lock_path.read_text(encoding="utf-8").splitlines()[0]
|
|
process_id = int(first_line.removeprefix("pid="))
|
|
os.kill(process_id, 0)
|
|
except PermissionError:
|
|
return True
|
|
except (IndexError, OSError, ValueError):
|
|
return False
|
|
return True
|