feat: Add CacheProvider API for external distributed caching

[deepme987]

Summary

Adds a public CacheProvider API that enables external cache providers to integrate with ComfyUI's caching system. This allows distributed caching across multiple ComfyUI instances (e.g., Kubernetes pods) without monkey-patching internal methods.

Key changes:

• New comfy_execution/cache_provider.py module with public API:
  • CacheProvider abstract base class with on_lookup, on_store, should_cache methods
  • CacheContext and CacheValue dataclasses for type-safe data passing
  • Thread-safe provider registry (register_cache_provider, get_cache_providers, etc.)
  • Utility functions: serialize_cache_key, contains_nan, estimate_value_size
• Modified comfy_execution/caching.py to call registered providers on cache miss/store
• Modified execution.py to notify providers on prompt start/end
• Deterministic cache key hashing using canonicalize + JSON (not pickle) for cross-pod consistency

Why this matters:

• Replaces fragile monkey-patching approach with a stable, public API
• Enables Comfy Cloud (and other deployments) to share cached results across pods
• Reduces redundant computation for expensive nodes (KSampler, VAEDecode, etc.)

Test plan

• Added unit tests in tests-unit/execution_test/test_cache_provider.py
  • Tests for _canonicalize deterministic ordering
  • Tests for serialize_cache_key hash consistency
  • Tests for contains_nan detection
  • Tests for estimate_value_size
  • Tests for provider registry (register/unregister/clear)
• Verified deterministic hashing works across Python sessions
• Tested with Comfy Cloud's SidecarCacheProvider implementation

🤖 Generated with Claude Code

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds an asynchronous, pluggable distributed cache provider system and integrates it into caching and execution. Introduces CacheContext, CacheValue, and an abstract CacheProvider with async on_lookup/on_store plus lifecycle hooks; a thread-safe provider registry and utilities for deterministic canonicalization, cache-key hashing, NaN detection, and size estimation. Converts core cache APIs to async, propagates per-prompt IDs to caches, consults external providers on misses and warms local cache on hits, and notifies providers of prompt lifecycle events from the executor. Adds unit tests covering utilities and registry behavior.

Changes

Cohort / File(s)	Summary
Public API Exposure comfy_api/latest/__init__.py	Adds public Caching class under Types that re-exports CacheProvider, CacheContext, CacheValue, and registration helpers (register_provider, unregister_provider); updates __all__ and Types.
Cache Provider Framework comfy_execution/cache_provider.py	Adds CacheContext and CacheValue dataclasses, CacheProvider ABC with async on_lookup/on_store, optional should_cache and lifecycle hooks, provider registry (register_cache_provider/unregister_cache_provider), and internal utilities for deterministic canonicalization, key serialization (SHA-256 with pickle fallback), NaN detection, and size estimation. Exposes registry query/management helpers.
Caching Layer Integration comfy_execution/caching.py	Converts core cache APIs to async, adds _is_subcache and _current_prompt_id propagation, consults external providers on cache misses and warms local cache on provider hits, fires provider stores asynchronously (fire-and-forget), and adds validation/safeguards for external caching (e.g., NaN handling). Updates many internal call paths to async and adds local sync variants (get_local, set_local).
Execution Integration execution.py	Switches execution to await async cache operations for outputs and objects, adds _notify_prompt_lifecycle(self, event: str, prompt_id: str) on PromptExecutor, and invokes prompt start/end notifications in a finally block.
Graph internals comfy_execution/graph.py	Replaces some uses of main cache accessors with local variants (get_local/set_local) in ExecutionList to avoid external-path interactions for those internal flows.
Unit Tests tests-unit/execution_test/test_cache_provider.py	Adds tests covering canonicalization, deterministic key hashing, NaN detection, size estimation (with optional PyTorch checks), provider registry behavior (register/unregister/clear/duplicates), dataclass integrity, and a MockCacheProvider for interaction tests.

Sequence Diagram

sequenceDiagram
    participant Executor as Executor
    participant LocalCache as Local Cache
    participant Provider as Cache Provider(s)
    
    rect rgba(100, 150, 200, 0.5)
    Note over Executor,Provider: Cache Lookup Flow
    Executor->>LocalCache: await get(node_id)
    LocalCache->>LocalCache: Check local storage
    alt Cache Hit
        LocalCache-->>Executor: Return cached value
    else Cache Miss
        LocalCache->>Provider: await on_lookup(context)
        Provider-->>LocalCache: CacheValue or None
        alt Provider Hit
            LocalCache->>LocalCache: Warm local cache
            LocalCache-->>Executor: Return external value
        else Provider Miss
            LocalCache-->>Executor: Return None
        end
    end
    end
    
    rect rgba(150, 200, 100, 0.5)
    Note over Executor,Provider: Cache Store Flow
    Executor->>LocalCache: await set(node_id, value)
    LocalCache->>LocalCache: Store locally
    LocalCache->>Provider: _notify_providers_store(context, value)
    Provider->>Provider: await on_store(context, value)
    Note over Provider: Fire-and-forget
    LocalCache-->>Executor: Complete
    end
    
    rect rgba(200, 150, 100, 0.5)
    Note over Executor,Provider: Lifecycle Management
    Executor->>Provider: on_prompt_start(prompt_id)
    Note over Provider: Initialize for prompt
    Executor->>Executor: Execute nodes with caching
    Executor->>Provider: on_prompt_end(prompt_id)
    Note over Provider: Cleanup for prompt
    end

Loading

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 19.15% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately and concisely describes the primary change: introducing a public CacheProvider API for external distributed caching.
Description check	✅ Passed	The description is comprehensive and directly related to the changeset, explaining the purpose, key changes, and rationale for the new CacheProvider API.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 5

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@comfy_execution/cache_provider.py`:
- Around line 178-202: The helper functions are defined with leading underscores
(_get_cache_providers, _has_cache_providers, _clear_cache_providers) while
callers in the PR import non-underscored names, causing import/runtime errors;
add thin public wrappers named get_cache_providers(), has_cache_providers(), and
clear_cache_providers() that simply call and return the existing underscored
implementations (preserving behavior around _providers_snapshot, _providers_lock
and _providers) or alternatively rename all callers to use the existing
underscored names—ensure the public wrappers delegate to _get_cache_providers,
_has_cache_providers, and _clear_cache_providers so existing internal
locking/caching behavior is preserved.
- Around line 260-267: Replace the process-local id()-based fallback in
comfy_execution/cache_provider.py so that when both deterministic serialization
and pickle serialization of cache_key fail the function returns None
(fail-closed) instead of hashing id(cache_key); specifically modify the except
block around pickle.dumps to return None on exception rather than
hashlib.sha256(str(id(cache_key)).encode()).hexdigest(). Then update
comfy_execution/caching.py to check the cache_key_hash variable and skip calling
the provider (do not attempt get/set on the provider) when cache_key_hash is
None so caching is bypassed for non-serializable keys.

In `@comfy_execution/caching.py`:
- Around line 231-235: Remove the unused imported names from the provider-hook
import blocks: in the import that currently imports _has_cache_providers,
_get_cache_providers, CacheContext, CacheValue, _serialize_cache_key,
_contains_nan, _logger remove CacheContext and _serialize_cache_key so only
_has_cache_providers, _get_cache_providers, CacheValue, _contains_nan, and
_logger remain; likewise remove CacheContext from the later import block (the
one around line ~270) so it does not get imported where unused. Keep the other
symbols (_has_cache_providers, _get_cache_providers, CacheValue, _contains_nan,
_logger) intact and rerun CI/lint to confirm F401 is resolved.
- Around line 212-223: The code calls self._check_providers_lookup on every
local cache miss (in _get_immediate), causing unnecessary external lookups for
non-output caches like caches.objects; restrict external provider checks to
output-value keys only by guarding the call with a check against the cache-key
type from cache_key_set (e.g., use a method such as
cache_key_set.is_output_key(node_id) or cache_key_set.is_output(cache_key) if
available) so that _check_providers_lookup(cache_key) is invoked only when the
key represents an output value; apply the same guard to the other affected block
(lines ~268-305) that currently triggers provider lookups on non-output misses.

In `@tests-unit/execution_test/test_cache_provider.py`:
- Around line 135-174: Tests assume serialize_cache_key returns raw SHA256 bytes
but the cache-provider contract returns a hex string; update assertions in
test_returns_bytes and test_dict_in_cache_key (and similar locations) to assert
isinstance(result, str) and len(result) == 64 (SHA256 hex length). Replace any
uses of non-existent CacheContext fields cache_key or cache_key_bytes with
cache_key_hash when constructing CacheContext instances (refer to CacheContext).
Finally, update tests that mock provider methods to follow the provider's async
signatures (use async def test helpers and await the provider methods or return
awaitable mocks) so they match the async method shape used by the provider
interface.

---

ℹ️ Review info

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 4f5bd39 and 4cbe4fe.

📒 Files selected for processing (5)

• comfy_api/latest/__init__.py
• comfy_execution/cache_provider.py
• comfy_execution/caching.py
• execution.py
• tests-unit/execution_test/test_cache_provider.py

[guill]

We should define these as small wrappers here rather than just importing them. As-is, someone could add/remove a param from register_provider and not realize that it'll be breaking the public API.

[deepme987]

something like this?

[guill]

Please follow the pattern of class Execution or class NodeReplacement in this file. To work with both process isolation and our versioning system:

1. This class must inherit from ProxiedSingleton.
2. The functions should be defined with this as a singleton (i.e. not staticmethods).
3. The functions should be defined as async. (People can get an automatically converted non-async version by importing ComfyAPISync.)

[guill]

nit: Should this be a debug-level log line?

[deepme987]

done

[guill]

Why do we need the prompt_id in the CacheContext? It seems like we might be able to simplify some things (e.g. storing the prompt id in the caches themselves) if we remove this. It also seems like we wouldn't want to care about the prompt id when determining caching (since the whole point of caching is to cache across different prompt ids).