From 38910648405fefd8faf0384e03a0188d69ae5f02 Mon Sep 17 00:00:00 2001 From: Deep Mehta Date: Thu, 12 Mar 2026 15:29:41 -0700 Subject: [PATCH] fix: address review feedback on cache provider API - Hold references to pending store tasks to prevent "Task was destroyed but it is still pending" warnings (bigcat88) - Parallel cache lookups with asyncio.gather instead of sequential awaits for better performance (bigcat88) - Delegate Caching.register/unregister_provider to existing functions in cache_provider.py instead of reimplementing (bigcat88) Co-Authored-By: Claude Opus 4.6 --- comfy_api/latest/__init__.py | 20 ++++---------------- comfy_execution/caching.py | 5 ++++- execution.py | 12 ++++++++---- 3 files changed, 16 insertions(+), 21 deletions(-) diff --git a/comfy_api/latest/__init__.py b/comfy_api/latest/__init__.py index c80ddb61a..04973fea0 100644 --- a/comfy_api/latest/__init__.py +++ b/comfy_api/latest/__init__.py @@ -107,25 +107,13 @@ class ComfyAPI_latest(ComfyAPIBase): async def register_provider(self, provider: "ComfyAPI_latest.Caching.CacheProvider") -> None: """Register an external cache provider. Providers are called in registration order.""" - from comfy_execution import cache_provider as _cp - with _cp._providers_lock: - if provider in _cp._providers: - _cp._logger.warning(f"Provider {provider.__class__.__name__} already registered") - return - _cp._providers.append(provider) - _cp._providers_snapshot = tuple(_cp._providers) - _cp._logger.debug(f"Registered cache provider: {provider.__class__.__name__}") + from comfy_execution.cache_provider import register_cache_provider + register_cache_provider(provider) async def unregister_provider(self, provider: "ComfyAPI_latest.Caching.CacheProvider") -> None: """Unregister a previously registered cache provider.""" - from comfy_execution import cache_provider as _cp - with _cp._providers_lock: - try: - _cp._providers.remove(provider) - _cp._providers_snapshot = tuple(_cp._providers) - _cp._logger.debug(f"Unregistered cache provider: {provider.__class__.__name__}") - except ValueError: - _cp._logger.warning(f"Provider {provider.__class__.__name__} was not registered") + from comfy_execution.cache_provider import unregister_cache_provider + unregister_cache_provider(provider) class ComfyExtension(ABC): async def on_load(self) -> None: diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index c5782d44a..750bddf2e 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -155,6 +155,7 @@ class BasicCache: self.cache_key_set: CacheKeySet self.cache = {} self.subcaches = {} + self._pending_store_tasks: set = set() async def set_prompt(self, dynprompt, node_ids, is_changed_cache): self.dynprompt = dynprompt @@ -253,7 +254,9 @@ class BasicCache: for provider in _get_cache_providers(): try: if provider.should_cache(context, cache_value): - asyncio.create_task(self._safe_provider_store(provider, context, cache_value)) + task = asyncio.create_task(self._safe_provider_store(provider, context, cache_value)) + self._pending_store_tasks.add(task) + task.add_done_callback(self._pending_store_tasks.discard) except Exception as e: _logger.warning(f"Cache provider {provider.__class__.__name__} error on store: {e}") diff --git a/execution.py b/execution.py index d4112b207..a8e8fc59f 100644 --- a/execution.py +++ b/execution.py @@ -726,10 +726,14 @@ class PromptExecutor: await cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache) cache.clean_unused() - cached_nodes = [] - for node_id in prompt: - if await self.caches.outputs.get(node_id) is not None: - cached_nodes.append(node_id) + node_ids = list(prompt.keys()) + cache_results = await asyncio.gather( + *(self.caches.outputs.get(node_id) for node_id in node_ids) + ) + cached_nodes = [ + node_id for node_id, result in zip(node_ids, cache_results) + if result is not None + ] comfy.model_management.cleanup_models_gc() self.add_message("execution_cached",