115 lines
3.9 KiB
Python
115 lines
3.9 KiB
Python
"""CPU-only native Comfy execution across private-slot lifetimes."""
|
|
import asyncio
|
|
import importlib
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import types
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
sys.path.insert(0, sys.argv[1])
|
|
import comfy.cli_args
|
|
comfy.cli_args.args.cpu = True
|
|
comfy.cli_args.args.disable_assets = True
|
|
import execution
|
|
import nodes
|
|
|
|
pkg = '_local_retention_probe'
|
|
for name, path in [(pkg, ROOT), (pkg + '.image_list', ROOT / 'image_list')]:
|
|
module = types.ModuleType(name)
|
|
module.__path__ = [str(path)]
|
|
sys.modules[name] = module
|
|
local = importlib.import_module(pkg + '.image_list.local_executor')
|
|
|
|
CALLS = []
|
|
OUTPUTS = []
|
|
|
|
class Counted:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {'required': {'value': ('INT',)}}
|
|
RETURN_TYPES = ('INT',)
|
|
FUNCTION = 'run'
|
|
def run(self, value):
|
|
if value < 0:
|
|
raise ValueError('intentional retention failure')
|
|
CALLS.append(value)
|
|
return (value,)
|
|
|
|
class Uncached(Counted):
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float('nan')
|
|
|
|
class Capture:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {'required': {'value': ('INT',), 'delivery': ('STRING',)}}
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = 'run'
|
|
def run(self, value, delivery):
|
|
OUTPUTS.append((value, delivery))
|
|
return ()
|
|
|
|
nodes.NODE_CLASS_MAPPINGS.update(RetentionCounted=Counted, RetentionUncached=Uncached, RetentionCapture=Capture)
|
|
|
|
async def run(slot, value, delivery, uncached=False):
|
|
prompt = {
|
|
'source': {'class_type': 'RetentionUncached' if uncached else 'RetentionCounted', 'inputs': {'value': value}},
|
|
'end': {'class_type': 'RetentionCapture', 'inputs': {'value': ['source', 0], 'delivery': delivery}},
|
|
}
|
|
await asyncio.wrap_future(slot.submit(prompt, delivery, 'end'))
|
|
|
|
async def main():
|
|
first = local.LocalExecutionSlot()
|
|
await run(first, 1, 'first-1')
|
|
await run(first, 2, 'first-2')
|
|
await first.close()
|
|
assert local._RETAINED_OUTPUTS is not None
|
|
second = local.LocalExecutionSlot()
|
|
await run(second, 1, 'second-1')
|
|
await run(second, 2, 'second-2')
|
|
await run(second, 3, 'changed')
|
|
await second.close()
|
|
assert CALLS == [1, 2, 3], CALLS
|
|
assert OUTPUTS == [(1, 'first-1'), (2, 'first-2'), (1, 'second-1'), (2, 'second-2'), (3, 'changed')], OUTPUTS
|
|
third = local.LocalExecutionSlot()
|
|
await run(third, 4, 'uncached-1', True)
|
|
await run(third, 4, 'uncached-2', True)
|
|
await third.close()
|
|
assert CALLS[-2:] == [4, 4]
|
|
comfy.cli_args.args.cache_none = True
|
|
disabled = local.LocalExecutionSlot()
|
|
await run(disabled, 1, 'disabled-1')
|
|
await run(disabled, 1, 'disabled-2')
|
|
await disabled.close()
|
|
assert CALLS[-2:] == [1, 1]
|
|
assert local._RETAINED_OUTPUTS is None
|
|
comfy.cli_args.args.cache_none = False
|
|
failed = local.LocalExecutionSlot()
|
|
await run(failed, 7, 'before-failure')
|
|
try:
|
|
await run(failed, -1, 'failure')
|
|
except RuntimeError:
|
|
pass
|
|
else:
|
|
raise AssertionError('failure fixture unexpectedly succeeded')
|
|
await failed.close()
|
|
assert local._RETAINED_OUTPUTS is None
|
|
cancelled = local.LocalExecutionSlot()
|
|
await run(cancelled, 8, 'cancelled')
|
|
cancelled.request_stop()
|
|
await cancelled.close()
|
|
assert local._RETAINED_OUTPUTS is None
|
|
oversized = execution.LRUCache(execution.CacheKeySetInputSignature, max_size=local.NATIVE_CACHE_MAX_ENTRIES)
|
|
empty = execution.DynamicPrompt({})
|
|
await oversized.set_prompt(empty, [], execution.IsChangedCache('bound', empty, oversized))
|
|
oversized.min_generation = oversized.generation
|
|
oversized.cache = {n: n for n in range(local.NATIVE_CACHE_MAX_ENTRIES + 1)}
|
|
local._retain_outputs(oversized)
|
|
assert local._RETAINED_OUTPUTS is None
|
|
print('PASS: cross-run native cache reuse, fresh delivery, input change, noncacheable nodes, cache_none, cancellation, retention bound')
|
|
|
|
asyncio.run(main())
|