mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-09-21 21:47:57 -05:00
refactor(assets): extract prompt_worker so its resume contract is testable in-process
This commit is contained in:
Simon Pinfold
@@ -0,0 +1,112 @@
|
||||
import gc
|
||||
import logging
|
||||
import time
|
||||
|
||||
import comfy.model_management
|
||||
import execution
|
||||
import hook_breaker_ac10a0
|
||||
from comfy.cli_args import args
|
||||
|
||||
|
||||
def prompt_worker(q, server_instance, asset_manager):
|
||||
current_time: float = 0.0
|
||||
cache_ram = 0
|
||||
cache_ram_inactive = 0
|
||||
if not args.cache_classic and not args.cache_none and args.cache_lru <= 0:
|
||||
cache_ram = min(10.0, max(2.0, comfy.model_management.total_ram * 0.10 / 1024.0))
|
||||
cache_ram_inactive = min(128.0, comfy.model_management.total_ram / 1024.0)
|
||||
if len(args.cache_ram) > 0:
|
||||
cache_ram = args.cache_ram[0]
|
||||
if len(args.cache_ram) > 1:
|
||||
cache_ram_inactive = args.cache_ram[1]
|
||||
|
||||
cache_type = execution.CacheType.RAM_PRESSURE
|
||||
if args.cache_classic:
|
||||
cache_type = execution.CacheType.CLASSIC
|
||||
elif args.cache_lru > 0:
|
||||
cache_type = execution.CacheType.LRU
|
||||
elif args.cache_none:
|
||||
cache_type = execution.CacheType.NONE
|
||||
|
||||
e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive }, asset_manager=asset_manager )
|
||||
last_gc_collect = 0
|
||||
need_gc = False
|
||||
gc_collect_interval = 10.0
|
||||
|
||||
while True:
|
||||
background_scan_paused = False
|
||||
try:
|
||||
timeout = 1000.0
|
||||
if need_gc:
|
||||
timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)
|
||||
|
||||
queue_item = q.get(timeout=timeout)
|
||||
if queue_item is not None:
|
||||
item, item_id = queue_item
|
||||
execution_start_time = time.perf_counter()
|
||||
prompt_id = item[1]
|
||||
server_instance.last_prompt_id = prompt_id
|
||||
|
||||
sensitive = item[5]
|
||||
extra_data = item[3].copy()
|
||||
for k in sensitive:
|
||||
extra_data[k] = sensitive[k]
|
||||
|
||||
asset_manager.pause_background_scan()
|
||||
background_scan_paused = True
|
||||
e.execute(item[2], prompt_id, extra_data, item[4])
|
||||
|
||||
need_gc = True
|
||||
|
||||
remove_sensitive = lambda prompt: prompt[:5] + prompt[6:]
|
||||
q.task_done(item_id,
|
||||
e.history_result,
|
||||
status=execution.PromptQueue.ExecutionStatus(
|
||||
status_str='success' if e.success else 'error',
|
||||
completed=e.success,
|
||||
messages=e.status_messages), process_item=remove_sensitive)
|
||||
if server_instance.client_id is not None:
|
||||
server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id)
|
||||
|
||||
current_time = time.perf_counter()
|
||||
execution_time = current_time - execution_start_time
|
||||
|
||||
# Log Time in a more readable way after 10 minutes
|
||||
if execution_time > 600:
|
||||
execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time))
|
||||
logging.info(f"Prompt executed in {execution_time}", extra={'color': 'green'})
|
||||
else:
|
||||
logging.info("Prompt executed in {:.2f} seconds".format(execution_time), extra={'color': 'green'})
|
||||
|
||||
flags = q.get_flags()
|
||||
free_memory = flags.get("free_memory", False)
|
||||
|
||||
if flags.get("unload_models", free_memory):
|
||||
comfy.model_management.unload_all_models()
|
||||
need_gc = True
|
||||
last_gc_collect = 0
|
||||
|
||||
if free_memory:
|
||||
e.reset()
|
||||
need_gc = True
|
||||
last_gc_collect = 0
|
||||
|
||||
if need_gc:
|
||||
current_time = time.perf_counter()
|
||||
if (current_time - last_gc_collect) > gc_collect_interval:
|
||||
gc.collect()
|
||||
comfy.model_management.soft_empty_cache()
|
||||
last_gc_collect = current_time
|
||||
need_gc = False
|
||||
hook_breaker_ac10a0.restore_functions()
|
||||
|
||||
asset_manager.queue_output_scan()
|
||||
asset_manager.resume_background_scan()
|
||||
background_scan_paused = False
|
||||
except BaseException:
|
||||
if background_scan_paused:
|
||||
try:
|
||||
asset_manager.resume_background_scan()
|
||||
except Exception:
|
||||
logging.exception("Failed to resume background asset scanning after prompt worker failure")
|
||||
raise
|
||||
@@ -240,7 +240,6 @@ execute_prestartup_script()
|
||||
# Main code
|
||||
import asyncio
|
||||
import threading
|
||||
import gc
|
||||
|
||||
if 'torch' in sys.modules:
|
||||
logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
|
||||
@@ -248,7 +247,7 @@ if 'torch' in sys.modules:
|
||||
|
||||
import comfy.utils
|
||||
|
||||
import execution
|
||||
from app.prompt_worker import prompt_worker
|
||||
import server
|
||||
from protocol import BinaryEventTypes
|
||||
import nodes
|
||||
@@ -316,110 +315,6 @@ def cuda_malloc_warning():
|
||||
logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")
|
||||
|
||||
|
||||
def prompt_worker(q, server_instance, asset_manager):
|
||||
current_time: float = 0.0
|
||||
cache_ram = 0
|
||||
cache_ram_inactive = 0
|
||||
if not args.cache_classic and not args.cache_none and args.cache_lru <= 0:
|
||||
cache_ram = min(10.0, max(2.0, comfy.model_management.total_ram * 0.10 / 1024.0))
|
||||
cache_ram_inactive = min(128.0, comfy.model_management.total_ram / 1024.0)
|
||||
if len(args.cache_ram) > 0:
|
||||
cache_ram = args.cache_ram[0]
|
||||
if len(args.cache_ram) > 1:
|
||||
cache_ram_inactive = args.cache_ram[1]
|
||||
|
||||
cache_type = execution.CacheType.RAM_PRESSURE
|
||||
if args.cache_classic:
|
||||
cache_type = execution.CacheType.CLASSIC
|
||||
elif args.cache_lru > 0:
|
||||
cache_type = execution.CacheType.LRU
|
||||
elif args.cache_none:
|
||||
cache_type = execution.CacheType.NONE
|
||||
|
||||
e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive }, asset_manager=asset_manager )
|
||||
last_gc_collect = 0
|
||||
need_gc = False
|
||||
gc_collect_interval = 10.0
|
||||
|
||||
while True:
|
||||
background_scan_paused = False
|
||||
try:
|
||||
timeout = 1000.0
|
||||
if need_gc:
|
||||
timeout = max(gc_collect_interval - (current_time - last_gc_collect), 0.0)
|
||||
|
||||
queue_item = q.get(timeout=timeout)
|
||||
if queue_item is not None:
|
||||
item, item_id = queue_item
|
||||
execution_start_time = time.perf_counter()
|
||||
prompt_id = item[1]
|
||||
server_instance.last_prompt_id = prompt_id
|
||||
|
||||
sensitive = item[5]
|
||||
extra_data = item[3].copy()
|
||||
for k in sensitive:
|
||||
extra_data[k] = sensitive[k]
|
||||
|
||||
asset_manager.pause_background_scan()
|
||||
background_scan_paused = True
|
||||
e.execute(item[2], prompt_id, extra_data, item[4])
|
||||
|
||||
need_gc = True
|
||||
|
||||
remove_sensitive = lambda prompt: prompt[:5] + prompt[6:]
|
||||
q.task_done(item_id,
|
||||
e.history_result,
|
||||
status=execution.PromptQueue.ExecutionStatus(
|
||||
status_str='success' if e.success else 'error',
|
||||
completed=e.success,
|
||||
messages=e.status_messages), process_item=remove_sensitive)
|
||||
if server_instance.client_id is not None:
|
||||
server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id)
|
||||
|
||||
current_time = time.perf_counter()
|
||||
execution_time = current_time - execution_start_time
|
||||
|
||||
# Log Time in a more readable way after 10 minutes
|
||||
if execution_time > 600:
|
||||
execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time))
|
||||
logging.info(f"Prompt executed in {execution_time}", extra={'color': 'green'})
|
||||
else:
|
||||
logging.info("Prompt executed in {:.2f} seconds".format(execution_time), extra={'color': 'green'})
|
||||
|
||||
flags = q.get_flags()
|
||||
free_memory = flags.get("free_memory", False)
|
||||
|
||||
if flags.get("unload_models", free_memory):
|
||||
comfy.model_management.unload_all_models()
|
||||
need_gc = True
|
||||
last_gc_collect = 0
|
||||
|
||||
if free_memory:
|
||||
e.reset()
|
||||
need_gc = True
|
||||
last_gc_collect = 0
|
||||
|
||||
if need_gc:
|
||||
current_time = time.perf_counter()
|
||||
if (current_time - last_gc_collect) > gc_collect_interval:
|
||||
gc.collect()
|
||||
comfy.model_management.soft_empty_cache()
|
||||
last_gc_collect = current_time
|
||||
need_gc = False
|
||||
hook_breaker_ac10a0.restore_functions()
|
||||
|
||||
asset_manager.queue_output_scan()
|
||||
asset_manager.resume_background_scan()
|
||||
background_scan_paused = False
|
||||
except BaseException:
|
||||
if background_scan_paused:
|
||||
try:
|
||||
asset_manager.resume_background_scan()
|
||||
except Exception:
|
||||
logging.exception("Failed to resume background asset scanning after prompt worker failure")
|
||||
raise
|
||||
|
||||
|
||||
async def run(server_instance, address='', port=8188, verbose=True, call_on_start=None):
|
||||
addresses = []
|
||||
for addr in address.split(","):
|
||||
|
||||
@@ -1,170 +1,102 @@
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import importlib
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _run_prompt_worker(script: str) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
cwd=Path(__file__).parents[2],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
class LoopEscape(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def test_prompt_worker_resumes_background_scan_when_execute_raises() -> None:
|
||||
script = """
|
||||
import sys
|
||||
|
||||
sys.argv = ["main.py", "--cpu"]
|
||||
|
||||
import main
|
||||
|
||||
class Queue:
|
||||
def __init__(self, completion_error: RuntimeError | None = None) -> None:
|
||||
self.completion_error = completion_error
|
||||
self.get_calls = 0
|
||||
|
||||
def get(self, timeout=None):
|
||||
self.get_calls += 1
|
||||
if self.get_calls > 1:
|
||||
raise LoopEscape("prompt worker requested a second item")
|
||||
return (0, "prompt-id", {}, {}, [], {}), 1
|
||||
|
||||
def task_done(self, *args, **kwargs) -> None:
|
||||
if self.completion_error is not None:
|
||||
raise self.completion_error
|
||||
|
||||
|
||||
class Server:
|
||||
last_prompt_id = None
|
||||
client_id = None
|
||||
|
||||
class BackgroundScan:
|
||||
paused = False
|
||||
|
||||
def pause_background_scan(self):
|
||||
class AssetManager:
|
||||
def __init__(self, resume_error: RuntimeError | None = None) -> None:
|
||||
self.paused = False
|
||||
self.resume_error = resume_error
|
||||
|
||||
def pause_background_scan(self) -> None:
|
||||
self.paused = True
|
||||
|
||||
def resume_background_scan(self):
|
||||
def resume_background_scan(self) -> None:
|
||||
self.paused = False
|
||||
if self.resume_error is not None:
|
||||
raise self.resume_error
|
||||
|
||||
|
||||
class Executor:
|
||||
def __init__(self, *args, **kwargs):
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
self.history_result = {}
|
||||
self.success = True
|
||||
self.status_messages = []
|
||||
|
||||
def execute(self, *args, **kwargs):
|
||||
raise RuntimeError("forced execute failure")
|
||||
|
||||
main.args.cache_classic = True
|
||||
main.execution.PromptExecutor = Executor
|
||||
asset_manager = BackgroundScan()
|
||||
|
||||
try:
|
||||
main.prompt_worker(Queue(), Server(), asset_manager)
|
||||
except RuntimeError as error:
|
||||
assert str(error) == "forced execute failure"
|
||||
|
||||
assert asset_manager.paused is False
|
||||
"""
|
||||
|
||||
result = _run_prompt_worker(script)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_prompt_worker_resumes_background_scan_when_completion_raises() -> None:
|
||||
script = """
|
||||
import sys
|
||||
|
||||
sys.argv = ["main.py", "--cpu"]
|
||||
|
||||
import main
|
||||
|
||||
class Queue:
|
||||
def get(self, timeout=None):
|
||||
return (0, "prompt-id", {}, {}, [], {}), 1
|
||||
|
||||
def task_done(self, *args, **kwargs):
|
||||
raise RuntimeError("forced completion failure")
|
||||
|
||||
class Server:
|
||||
last_prompt_id = None
|
||||
client_id = None
|
||||
|
||||
class BackgroundScan:
|
||||
paused = False
|
||||
|
||||
def pause_background_scan(self):
|
||||
self.paused = True
|
||||
|
||||
def resume_background_scan(self):
|
||||
self.paused = False
|
||||
|
||||
class Executor:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.history_result = {}
|
||||
self.success = True
|
||||
self.status_messages = []
|
||||
|
||||
def execute(self, *args, **kwargs):
|
||||
def execute(self, *args, **kwargs) -> None:
|
||||
return None
|
||||
|
||||
main.args.cache_classic = True
|
||||
main.execution.PromptExecutor = Executor
|
||||
asset_manager = BackgroundScan()
|
||||
|
||||
try:
|
||||
main.prompt_worker(Queue(), Server(), asset_manager)
|
||||
except RuntimeError as error:
|
||||
assert str(error) == "forced completion failure"
|
||||
|
||||
assert asset_manager.paused is False
|
||||
"""
|
||||
|
||||
result = _run_prompt_worker(script)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_prompt_worker_preserves_execute_error_when_resume_raises() -> None:
|
||||
script = """
|
||||
import sys
|
||||
|
||||
sys.argv = ["main.py", "--cpu"]
|
||||
|
||||
import main
|
||||
|
||||
class Queue:
|
||||
def get(self, timeout=None):
|
||||
return (0, "prompt-id", {}, {}, [], {}), 1
|
||||
|
||||
class Server:
|
||||
last_prompt_id = None
|
||||
client_id = None
|
||||
|
||||
class BackgroundScan:
|
||||
paused = False
|
||||
|
||||
def pause_background_scan(self):
|
||||
self.paused = True
|
||||
|
||||
def resume_background_scan(self):
|
||||
self.paused = False
|
||||
raise RuntimeError("forced resume failure")
|
||||
|
||||
class Executor:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def execute(self, *args, **kwargs):
|
||||
class ExecuteFailureExecutor(Executor):
|
||||
def execute(self, *args, **kwargs) -> None:
|
||||
raise RuntimeError("forced execute failure")
|
||||
|
||||
main.args.cache_classic = True
|
||||
main.execution.PromptExecutor = Executor
|
||||
asset_manager = BackgroundScan()
|
||||
|
||||
try:
|
||||
main.prompt_worker(Queue(), Server(), asset_manager)
|
||||
except RuntimeError as error:
|
||||
assert str(error) == "forced execute failure"
|
||||
else:
|
||||
raise AssertionError("prompt_worker should propagate the execute failure")
|
||||
@pytest.fixture
|
||||
def prompt_worker_module(monkeypatch):
|
||||
from comfy.cli_args import args
|
||||
|
||||
assert asset_manager.paused is False
|
||||
"""
|
||||
monkeypatch.setattr(args, "cpu", True, raising=False)
|
||||
try:
|
||||
return importlib.import_module("app.prompt_worker")
|
||||
except Exception as exc:
|
||||
pytest.skip(f"prompt worker module could not be imported in CPU mode: {exc!r}")
|
||||
|
||||
result = _run_prompt_worker(script)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
def test_prompt_worker_resumes_background_scan_when_execute_raises(prompt_worker_module, monkeypatch) -> None:
|
||||
monkeypatch.setattr(prompt_worker_module.execution, "PromptExecutor", ExecuteFailureExecutor)
|
||||
asset_manager = AssetManager()
|
||||
|
||||
with pytest.raises(RuntimeError, match="^forced execute failure$"):
|
||||
prompt_worker_module.prompt_worker(Queue(), Server(), asset_manager)
|
||||
|
||||
assert asset_manager.paused is False
|
||||
|
||||
|
||||
def test_prompt_worker_resumes_background_scan_when_completion_raises(prompt_worker_module, monkeypatch) -> None:
|
||||
monkeypatch.setattr(prompt_worker_module.execution, "PromptExecutor", Executor)
|
||||
asset_manager = AssetManager()
|
||||
|
||||
with pytest.raises(RuntimeError, match="^forced completion failure$"):
|
||||
prompt_worker_module.prompt_worker(
|
||||
Queue(completion_error=RuntimeError("forced completion failure")),
|
||||
Server(),
|
||||
asset_manager,
|
||||
)
|
||||
|
||||
assert asset_manager.paused is False
|
||||
|
||||
|
||||
def test_prompt_worker_preserves_execute_error_when_resume_raises(prompt_worker_module, monkeypatch) -> None:
|
||||
monkeypatch.setattr(prompt_worker_module.execution, "PromptExecutor", ExecuteFailureExecutor)
|
||||
asset_manager = AssetManager(resume_error=RuntimeError("forced resume failure"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="^forced execute failure$"):
|
||||
prompt_worker_module.prompt_worker(Queue(), Server(), asset_manager)
|
||||
|
||||
assert asset_manager.paused is False
|
||||
|
||||
Reference in New Issue
Block a user