Ken SimpsonandClaude Opus 5.5
62f0a13acc
Current ComfyUI's PromptExecutor.execute_async indexes cache_args["ram"] and cache_args["ram_inactive"] on every run. Local fallback execution built its nested executor without cache_args, so every locally executed subgraph failed with "'NoneType' object is not subscriptable". Pass zeroed classic-cache args, and restore ComfyUI's module-global RAM cache release state after the nested run so the host executor keeps its RAM-pressure eviction. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
106 lines
3.7 KiB
Python
106 lines
3.7 KiB
Python
"""Tests for the nested PromptExecutor used by local fallback execution."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from types import ModuleType, SimpleNamespace
|
|
from typing import Any
|
|
|
|
|
|
def _install_fake_memory_management(monkeypatch: Any) -> ModuleType:
|
|
"""Install a minimal stand-in for ComfyUI's RAM cache release globals."""
|
|
memory_management = ModuleType("comfy.memory_management")
|
|
memory_management.extra_ram_release_callback = None
|
|
memory_management.RAM_CACHE_HEADROOM = 0
|
|
|
|
def set_ram_cache_release_state(callback: Any, headroom: int) -> None:
|
|
"""Mirror ComfyUI's module-global setter."""
|
|
memory_management.extra_ram_release_callback = callback
|
|
memory_management.RAM_CACHE_HEADROOM = max(0, int(headroom))
|
|
|
|
memory_management.set_ram_cache_release_state = set_ram_cache_release_state
|
|
monkeypatch.setitem(sys.modules, "comfy.memory_management", memory_management)
|
|
return memory_management
|
|
|
|
|
|
def _fake_execution_module(memory_management: ModuleType) -> SimpleNamespace:
|
|
"""Return an execution module whose executor behaves like current ComfyUI."""
|
|
|
|
class FakePromptExecutor:
|
|
"""Require RAM cache args and clobber the global release state like ComfyUI."""
|
|
|
|
def __init__(
|
|
self,
|
|
server: Any,
|
|
cache_type: Any = False,
|
|
cache_args: Any = None,
|
|
) -> None:
|
|
"""Record construction arguments."""
|
|
self.server = server
|
|
self.cache_type = cache_type
|
|
self.cache_args = cache_args
|
|
self.executed_outputs: list[str] = []
|
|
|
|
def execute(
|
|
self,
|
|
*,
|
|
prompt: dict[str, Any],
|
|
prompt_id: str,
|
|
extra_data: dict[str, Any],
|
|
execute_outputs: list[str],
|
|
) -> None:
|
|
"""Fail like ComfyUI on missing cache args, then reset global state."""
|
|
del prompt, prompt_id, extra_data
|
|
headroom = int(self.cache_args["ram"] * (1024**3))
|
|
float(self.cache_args["ram_inactive"])
|
|
memory_management.set_ram_cache_release_state(None, headroom)
|
|
self.executed_outputs = execute_outputs
|
|
|
|
return SimpleNamespace(PromptExecutor=FakePromptExecutor)
|
|
|
|
|
|
def test_nested_executor_supplies_ram_cache_args(
|
|
local_execution_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""Current ComfyUI indexes cache_args['ram'] on every execute."""
|
|
memory_management = _install_fake_memory_management(monkeypatch)
|
|
|
|
executor = local_execution_module._run_nested_prompt_executor(
|
|
_fake_execution_module(memory_management),
|
|
prompt={},
|
|
prompt_id="prompt-1",
|
|
extra_data={},
|
|
execute_outputs=["9"],
|
|
)
|
|
|
|
assert executor.cache_type is False
|
|
assert executor.cache_args == {"lru": 0, "ram": 0.0, "ram_inactive": 0.0}
|
|
assert executor.executed_outputs == ["9"]
|
|
|
|
|
|
def test_nested_executor_restores_host_ram_release_state(
|
|
local_execution_module: Any,
|
|
monkeypatch: Any,
|
|
) -> None:
|
|
"""The host executor's RAM-pressure eviction must survive a nested run."""
|
|
memory_management = _install_fake_memory_management(monkeypatch)
|
|
|
|
def host_release(target: int, free_active: bool = False) -> int:
|
|
"""Stand in for the host cache's RAM release callback."""
|
|
del target, free_active
|
|
return 0
|
|
|
|
memory_management.set_ram_cache_release_state(host_release, 8 * 1024**3)
|
|
|
|
local_execution_module._run_nested_prompt_executor(
|
|
_fake_execution_module(memory_management),
|
|
prompt={},
|
|
prompt_id="prompt-1",
|
|
extra_data={},
|
|
execute_outputs=["9"],
|
|
)
|
|
|
|
assert memory_management.extra_ram_release_callback is host_release
|
|
assert memory_management.RAM_CACHE_HEADROOM == 8 * 1024**3
|