mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-30 19:07:25 +08:00
fcbe7db46f
Adds a new on_set_prompt() lifecycle hook on CacheProvider that fires after the cache key set is prepared for a new prompt. Dispatched via asyncio.create_task with errors swallowed (same fail-safe pattern as on_store / on_lookup). Why: BasicCache's lifecycle notifications to external providers were incomplete. set_prompt is a key per-prompt event that providers need visibility into — for example, to reset per-prompt timing/state used for cost-aware caching policies (a provider can set t=0 here, then measure elapsed at each on_store to estimate compute saved by a hit). Backward-compatible: default implementation is a no-op, existing providers compile and run unchanged. Providers that need the per-prompt boundary override on_set_prompt().
47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import Optional
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass
|
|
class CacheContext:
|
|
node_id: str
|
|
class_type: str
|
|
cache_key_hash: str # SHA256 hex digest
|
|
|
|
|
|
@dataclass
|
|
class CacheValue:
|
|
outputs: list
|
|
ui: dict = None
|
|
|
|
|
|
class CacheProvider(ABC):
|
|
"""Abstract base class for external cache providers.
|
|
Exceptions from provider methods are caught by the caller and never break execution.
|
|
"""
|
|
|
|
async def on_set_prompt(self) -> None:
|
|
"""Called after prompt cache keys are prepared. Dispatched via asyncio.create_task."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def on_lookup(self, context: CacheContext) -> Optional[CacheValue]:
|
|
"""Called on local cache miss. Return CacheValue if found, None otherwise."""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def on_store(self, context: CacheContext, value: CacheValue) -> None:
|
|
"""Called after local store. Dispatched via asyncio.create_task."""
|
|
pass
|
|
|
|
def should_cache(self, context: CacheContext, value: Optional[CacheValue] = None) -> bool:
|
|
"""Return False to skip external caching for this node. Default: True."""
|
|
return True
|
|
|
|
def on_prompt_start(self, prompt_id: str) -> None:
|
|
pass
|
|
|
|
def on_prompt_end(self, prompt_id: str) -> None:
|
|
pass
|