477 lines
17 KiB
Python
477 lines
17 KiB
Python
"""Idle-gated, deduplicated Cloudflare R2 write-back coordination."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import OrderedDict
|
|
from contextlib import AbstractContextManager, nullcontext
|
|
from dataclasses import dataclass, field
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Any, Callable, Protocol
|
|
|
|
if __package__:
|
|
from .r2_cache import (
|
|
R2CacheClient,
|
|
R2CacheError,
|
|
R2DownloadRequest,
|
|
R2UploadPlan,
|
|
)
|
|
from .sync_protocols import (
|
|
CancellationCheck,
|
|
CancellableR2WriteBackBackend,
|
|
R2MaterializingBackend,
|
|
R2WorkerPreflightBackend,
|
|
SyncStatusCallback,
|
|
_ContentAddressedSyncResult,
|
|
_ContentAddressedSyncSpec,
|
|
_R2MaterializationOutcome,
|
|
_format_indexed_asset_status,
|
|
)
|
|
else: # pragma: no cover - flat import inside the Modal container.
|
|
from r2_cache import R2CacheClient, R2CacheError, R2DownloadRequest, R2UploadPlan
|
|
from sync_protocols import (
|
|
CancellationCheck,
|
|
CancellableR2WriteBackBackend,
|
|
R2MaterializingBackend,
|
|
R2WorkerPreflightBackend,
|
|
SyncStatusCallback,
|
|
_ContentAddressedSyncResult,
|
|
_ContentAddressedSyncSpec,
|
|
_R2MaterializationOutcome,
|
|
_format_indexed_asset_status,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class R2WriteBackCancelled(RuntimeError):
|
|
"""Signal that an idle write-back yielded to foreground workflow activity."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _R2WriteBackJob:
|
|
"""Describe one deduplicated best-effort cache population operation."""
|
|
|
|
key: tuple[str, ...]
|
|
callback: Callable[[CancellationCheck], None] = field(compare=False, repr=False)
|
|
|
|
|
|
class R2WriteBackCoordinator:
|
|
"""Run cache population only while no remote workflow is active."""
|
|
|
|
def __init__(self, *, max_workers: int = 2, max_pending_jobs: int = 1024) -> None:
|
|
"""Start bounded daemon workers for opportunistic write-back jobs."""
|
|
if max_workers <= 0 or max_pending_jobs <= 0:
|
|
raise ValueError("R2 write-back coordinator bounds must be positive.")
|
|
self._condition = threading.Condition()
|
|
self._pending: OrderedDict[tuple[str, ...], _R2WriteBackJob] = OrderedDict()
|
|
self._known_keys: set[tuple[str, ...]] = set()
|
|
self._active_prompts: set[str] = set()
|
|
self._active_cancellations: set[threading.Event] = set()
|
|
self._active_jobs = 0
|
|
self._max_pending_jobs = max_pending_jobs
|
|
self._workers = tuple(
|
|
threading.Thread(
|
|
target=self._worker,
|
|
name=f"comfy-r2-writeback-{worker_index + 1}",
|
|
daemon=True,
|
|
)
|
|
for worker_index in range(max_workers)
|
|
)
|
|
for worker in self._workers:
|
|
worker.start()
|
|
|
|
def begin_prompt(self, prompt_id: str) -> None:
|
|
"""Reserve foreground priority and preempt active cache transfers."""
|
|
normalized_prompt_id = str(prompt_id).strip()
|
|
if not normalized_prompt_id:
|
|
raise ValueError("R2 write-back prompt ID must not be empty.")
|
|
with self._condition:
|
|
self._active_prompts.add(normalized_prompt_id)
|
|
for cancellation in self._active_cancellations:
|
|
cancellation.set()
|
|
self._condition.notify_all()
|
|
|
|
def finish_prompt(self, prompt_id: str) -> None:
|
|
"""Release one prompt reservation and wake idle cache workers."""
|
|
normalized_prompt_id = str(prompt_id).strip()
|
|
if not normalized_prompt_id:
|
|
return
|
|
with self._condition:
|
|
self._active_prompts.discard(normalized_prompt_id)
|
|
self._condition.notify_all()
|
|
|
|
def reset_prompt_reservations_for_tests(self) -> None:
|
|
"""Clear synthetic prompt reservations left by isolated queue tests."""
|
|
with self._condition:
|
|
self._active_prompts.clear()
|
|
for cancellation in self._active_cancellations:
|
|
cancellation.set()
|
|
self._condition.notify_all()
|
|
|
|
def submit(
|
|
self,
|
|
key: tuple[str, ...],
|
|
callback: Callable[[CancellationCheck], None],
|
|
) -> bool:
|
|
"""Queue one deduplicated job without waiting for capacity or execution."""
|
|
if not key or any(not str(part) for part in key):
|
|
raise ValueError("R2 write-back job key must contain non-empty values.")
|
|
with self._condition:
|
|
if key in self._known_keys:
|
|
return False
|
|
if len(self._pending) >= self._max_pending_jobs:
|
|
logger.warning(
|
|
"Dropping R2 write-back because the idle queue reached %d jobs key=%s.",
|
|
self._max_pending_jobs,
|
|
key[-1],
|
|
)
|
|
return False
|
|
job = _R2WriteBackJob(key=key, callback=callback)
|
|
self._pending[key] = job
|
|
self._known_keys.add(key)
|
|
self._condition.notify()
|
|
return True
|
|
|
|
def wait_for_idle(self, timeout_seconds: float | None = None) -> bool:
|
|
"""Wait until pending and active jobs drain, primarily for tests and shutdown."""
|
|
deadline = (
|
|
None if timeout_seconds is None else time.monotonic() + timeout_seconds
|
|
)
|
|
with self._condition:
|
|
while self._pending or self._active_jobs:
|
|
if deadline is None:
|
|
self._condition.wait()
|
|
continue
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0.0:
|
|
return False
|
|
self._condition.wait(remaining)
|
|
return True
|
|
|
|
def _worker(self) -> None:
|
|
"""Run queued jobs only during foreground-idle windows."""
|
|
while True:
|
|
with self._condition:
|
|
self._condition.wait_for(
|
|
lambda: bool(self._pending) and not self._active_prompts
|
|
)
|
|
key, job = self._pending.popitem(last=False)
|
|
cancellation = threading.Event()
|
|
self._active_cancellations.add(cancellation)
|
|
self._active_jobs += 1
|
|
cancelled = False
|
|
try:
|
|
job.callback(cancellation.is_set)
|
|
except R2WriteBackCancelled:
|
|
cancelled = True
|
|
except (
|
|
AssertionError,
|
|
KeyError,
|
|
OSError,
|
|
RuntimeError,
|
|
TypeError,
|
|
ValueError,
|
|
) as exc:
|
|
logger.warning(
|
|
"Unhandled R2 write-back job failure key=%s: %s",
|
|
key[-1],
|
|
exc,
|
|
)
|
|
finally:
|
|
with self._condition:
|
|
self._active_cancellations.discard(cancellation)
|
|
self._active_jobs -= 1
|
|
if cancelled:
|
|
self._pending.setdefault(key, job)
|
|
else:
|
|
self._known_keys.discard(key)
|
|
self._condition.notify_all()
|
|
|
|
|
|
_R2_WRITE_BACK_COORDINATOR = R2WriteBackCoordinator()
|
|
|
|
|
|
def begin_r2_writeback_prompt(prompt_id: str) -> None:
|
|
"""Pause opportunistic cache population for one foreground prompt."""
|
|
_R2_WRITE_BACK_COORDINATOR.begin_prompt(prompt_id)
|
|
|
|
|
|
def finish_r2_writeback_prompt(prompt_id: str) -> None:
|
|
"""Resume cache population after one foreground prompt leaves the queue."""
|
|
_R2_WRITE_BACK_COORDINATOR.finish_prompt(prompt_id)
|
|
|
|
|
|
def _never_cancel() -> bool:
|
|
"""Return false for write-backs without a cooperative cancellation source."""
|
|
return False
|
|
|
|
|
|
def _emit_sync_status(
|
|
status_callback: SyncStatusCallback | None,
|
|
message: str,
|
|
current: int | None = None,
|
|
total: int | None = None,
|
|
) -> None:
|
|
"""Emit one synchronization status update when available."""
|
|
if status_callback is not None:
|
|
status_callback(message, current, total)
|
|
|
|
|
|
class R2TransferHost(Protocol):
|
|
"""Define engine services used by R2 transfer coordination."""
|
|
|
|
def _destination_label(self) -> str:
|
|
"""Return a user-facing worker destination."""
|
|
|
|
def _raise_if_cancelled(self) -> None:
|
|
"""Raise when foreground synchronization was cancelled."""
|
|
|
|
def _store_sync_record(
|
|
self,
|
|
*,
|
|
sync_key: str,
|
|
remote_path: str,
|
|
source_description: str,
|
|
) -> None:
|
|
"""Persist one successful synchronization record."""
|
|
|
|
|
|
@dataclass
|
|
class R2TransferManager:
|
|
"""Own R2 preflight, read-through, and background write-back behavior."""
|
|
|
|
_host: R2TransferHost
|
|
volume: Any
|
|
r2_cache: R2CacheClient | None = None
|
|
cancellation_check: CancellationCheck | None = None
|
|
r2_writeback_activity: Callable[[], AbstractContextManager[None]] | None = None
|
|
|
|
def preflight_r2_access(
|
|
self,
|
|
*,
|
|
status_callback: SyncStatusCallback | None = None,
|
|
) -> None:
|
|
"""Fail before asset transfer when a worker cannot use configured R2 URLs."""
|
|
self._host._raise_if_cancelled()
|
|
if self.r2_cache is None or not isinstance(
|
|
self.volume,
|
|
R2WorkerPreflightBackend,
|
|
):
|
|
return
|
|
destination = self._host._destination_label()
|
|
_emit_sync_status(
|
|
status_callback,
|
|
f"Checking Cloudflare R2 access from {destination}",
|
|
)
|
|
try:
|
|
request = self.r2_cache.worker_preflight_request()
|
|
except (R2CacheError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
|
raise R2CacheError(
|
|
"Cloudflare R2 controller validation failed before the worker "
|
|
f"preflight: {exc}"
|
|
) from exc
|
|
try:
|
|
self.volume.preflight_r2_access(
|
|
request,
|
|
cancellation_check=self.cancellation_check,
|
|
)
|
|
except InterruptedError as exc:
|
|
self._host._raise_if_cancelled()
|
|
raise R2CacheError(
|
|
f"Cloudflare R2 worker preflight was interrupted for {destination}."
|
|
) from exc
|
|
except (R2CacheError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
|
diagnostic = str(exc)
|
|
hint = (
|
|
" Remove Client IP Address Filtering from the bucket-scoped R2 "
|
|
"API token; dynamic worker egress addresses cannot be reliably "
|
|
"allowlisted."
|
|
if "r2_code=AccessDenied" in diagnostic
|
|
else ""
|
|
)
|
|
raise R2CacheError(
|
|
f"Cloudflare R2 is reachable from the controller but not from "
|
|
f"{destination}.{hint} Safe worker diagnostic: {diagnostic}"
|
|
) from exc
|
|
self._host._raise_if_cancelled()
|
|
logger.info(
|
|
"Validated Cloudflare R2 access from %s before asset transfer.",
|
|
destination,
|
|
)
|
|
|
|
def _materialize_r2_source(
|
|
self,
|
|
spec: _ContentAddressedSyncSpec,
|
|
size_bytes: int,
|
|
) -> _R2MaterializationOutcome:
|
|
"""Try one signed worker-side R2 download and preserve upload fallback."""
|
|
if self.r2_cache is None or not isinstance(self.volume, R2MaterializingBackend):
|
|
return _R2MaterializationOutcome(result=None)
|
|
request: R2DownloadRequest | None = None
|
|
try:
|
|
request = self.r2_cache.download_request(spec.sha256, size_bytes)
|
|
if request is None:
|
|
return _R2MaterializationOutcome(result=None)
|
|
_emit_sync_status(
|
|
spec.status_callback,
|
|
_format_indexed_asset_status(
|
|
spec.local_path.name,
|
|
action="Downloading asset",
|
|
location="from Cloudflare R2",
|
|
item_index=spec.status_current,
|
|
total_items=spec.status_total,
|
|
),
|
|
spec.status_current,
|
|
spec.status_total,
|
|
)
|
|
self.volume.materialize_r2_file(
|
|
request,
|
|
spec.remote_path,
|
|
cancellation_check=self.cancellation_check,
|
|
)
|
|
except InterruptedError as exc:
|
|
raise SyncCancelledError("Remote workflow preparation was cancelled.") from exc
|
|
except (R2CacheError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
|
logger.warning(
|
|
"Cloudflare R2 materialization failed for SHA-256 %s; "
|
|
"falling back to the existing upload path: %s",
|
|
spec.sha256,
|
|
exc,
|
|
)
|
|
return _R2MaterializationOutcome(
|
|
result=None,
|
|
refresh_required=request is not None,
|
|
)
|
|
self._host._store_sync_record(
|
|
sync_key=spec.sync_key,
|
|
remote_path=spec.remote_path,
|
|
source_description=f"Cloudflare R2 SHA-256 {spec.sha256}",
|
|
)
|
|
logger.info(
|
|
"Materialized SHA-256 %s from Cloudflare R2 at %s.",
|
|
spec.sha256,
|
|
spec.remote_path,
|
|
)
|
|
return _R2MaterializationOutcome(
|
|
result=_ContentAddressedSyncResult(
|
|
remote_path=spec.remote_path,
|
|
uploaded=True,
|
|
)
|
|
)
|
|
|
|
def _schedule_r2_writeback(
|
|
self,
|
|
*,
|
|
sha256: str,
|
|
size_bytes: int,
|
|
remote_path: str,
|
|
force: bool = False,
|
|
) -> None:
|
|
"""Queue one best-effort R2 write without delaying foreground progress."""
|
|
if not self._r2_writeback_enabled():
|
|
return
|
|
assert self.r2_cache is not None
|
|
job_key = (
|
|
*self._r2_writeback_job_prefix(),
|
|
sha256,
|
|
str(size_bytes),
|
|
remote_path,
|
|
"force" if force else "normal",
|
|
)
|
|
_R2_WRITE_BACK_COORDINATOR.submit(
|
|
job_key,
|
|
lambda cancellation_check: self._write_back_r2_file(
|
|
sha256,
|
|
size_bytes,
|
|
remote_path,
|
|
force,
|
|
cancellation_check=cancellation_check,
|
|
),
|
|
)
|
|
|
|
def _r2_writeback_job_prefix(self) -> tuple[str, str, str]:
|
|
"""Return a stable cache namespace used to deduplicate background jobs."""
|
|
if self.r2_cache is None:
|
|
raise RuntimeError("R2 write-back job requested without an R2 cache.")
|
|
configuration = getattr(self.r2_cache, "configuration", None)
|
|
if configuration is not None:
|
|
return (
|
|
str(configuration.endpoint_url),
|
|
str(configuration.bucket),
|
|
str(configuration.key_prefix),
|
|
)
|
|
return (
|
|
type(self.r2_cache).__module__,
|
|
type(self.r2_cache).__qualname__,
|
|
str(id(self.r2_cache)),
|
|
)
|
|
|
|
def _r2_writeback_enabled(self) -> bool:
|
|
"""Return whether this engine can populate its configured R2 cache."""
|
|
return bool(
|
|
self.r2_cache is not None
|
|
and self.r2_cache.write_back_mode != "off"
|
|
and isinstance(self.volume, R2MaterializingBackend)
|
|
)
|
|
|
|
def _write_back_r2_file(
|
|
self,
|
|
sha256: str,
|
|
size_bytes: int,
|
|
remote_path: str,
|
|
force: bool = False,
|
|
*,
|
|
cancellation_check: CancellationCheck = _never_cancel,
|
|
) -> None:
|
|
"""Upload one worker-resident file to R2 without exposing permanent keys."""
|
|
if self.r2_cache is None or not isinstance(
|
|
self.volume,
|
|
R2MaterializingBackend,
|
|
):
|
|
return
|
|
plan: R2UploadPlan | None = None
|
|
try:
|
|
if cancellation_check():
|
|
raise R2WriteBackCancelled("R2 write-back yielded before activity.")
|
|
activity = (
|
|
self.r2_writeback_activity()
|
|
if self.r2_writeback_activity is not None
|
|
else nullcontext()
|
|
)
|
|
with activity:
|
|
if cancellation_check():
|
|
raise R2WriteBackCancelled("R2 write-back yielded before signing.")
|
|
plan = self.r2_cache.prepare_upload(sha256, size_bytes, force=force)
|
|
if plan is None:
|
|
return
|
|
if cancellation_check():
|
|
raise R2WriteBackCancelled("R2 write-back yielded after signing.")
|
|
if isinstance(self.volume, CancellableR2WriteBackBackend):
|
|
result = self.volume.upload_r2_file_cancellable(
|
|
plan,
|
|
remote_path,
|
|
cancellation_check=cancellation_check,
|
|
)
|
|
else:
|
|
result = self.volume.upload_r2_file(plan, remote_path)
|
|
if cancellation_check():
|
|
raise R2WriteBackCancelled(
|
|
"R2 write-back yielded before completion."
|
|
)
|
|
self.r2_cache.complete_upload(plan, result)
|
|
logger.info("Wrote SHA-256 %s back to Cloudflare R2.", sha256)
|
|
except (InterruptedError, R2WriteBackCancelled) as exc:
|
|
if plan is not None:
|
|
self.r2_cache.abort_upload(plan)
|
|
raise R2WriteBackCancelled(str(exc)) from exc
|
|
except (R2CacheError, OSError, RuntimeError, TypeError, ValueError) as exc:
|
|
if plan is not None:
|
|
self.r2_cache.abort_upload(plan)
|
|
logger.warning("Cloudflare R2 write-back failed for SHA-256 %s: %s", sha256, exc)
|
|
|
|
def wait_for_r2_writebacks(self) -> None:
|
|
"""Wait for this engine's currently scheduled R2 writes to finish."""
|
|
_R2_WRITE_BACK_COORDINATOR.wait_for_idle()
|