From 05432b7c28d9e1a0fc68635706e1ed5277979afa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=8E=E7=94=9F=E3=81=AE=E7=94=B7?= Date: Mon, 20 Jul 2026 02:00:19 +0900 Subject: [PATCH] Fix continuous batching request isolation --- comfy/continuous_batching.py | 26 ++++- execution.py | 16 +-- main.py | 20 +++- nodes.py | 3 + .../comfy_test/test_continuous_batching.py | 67 ++++++++++-- .../execution_test/test_continuous_worker.py | 101 ++++++++++++++++++ 6 files changed, 213 insertions(+), 20 deletions(-) diff --git a/comfy/continuous_batching.py b/comfy/continuous_batching.py index 072f02d06..cb940d706 100644 --- a/comfy/continuous_batching.py +++ b/comfy/continuous_batching.py @@ -627,6 +627,10 @@ class ContinuousBatchCoordinator: for request, (state, finished) in zip(current, updates): if request.cancelled: state.clear() + elif _is_cancelled(getattr(state, "prompt_id", None)): + state.clear() + if not request.future.done(): + request.future.set_exception(comfy.model_management.InterruptProcessingException()) elif finished: finished_any = True output = state.output @@ -672,6 +676,9 @@ class ContinuousVAEDecodeRequest: max_batch_size: int admission_delay: float prompt_id: str | None = None + node_id: str | None = None + client_id: str | None = None + progress_registry: Any = None def validate(self): if type(self.vae) is not comfy.sd.VAE: @@ -720,6 +727,7 @@ class ContinuousVAEDecodeRequest: def clear(self): self.vae = None self.samples = None + self.progress_registry = None @dataclass @@ -786,10 +794,20 @@ class ContinuousVAECoordinator: return group def _decode(self, group): - if len(group) == 1: - images = self.vae.decode(group[0].state.samples) - else: - images = self.vae.decode(torch.cat([request.state.samples for request in group])) + representative = group[0].state + samples = representative.samples if len(group) == 1 else torch.cat([request.state.samples for request in group]) + client_id_token = set_current_client_id(representative.client_id) + progress_token = set_progress_registry(representative.progress_registry) if representative.progress_registry is not None else None + try: + if representative.prompt_id is not None and representative.node_id is not None: + with CurrentNodeContext(representative.prompt_id, representative.node_id): + images = self.vae.decode(samples) + else: + images = self.vae.decode(samples) + finally: + if progress_token is not None: + reset_progress_registry(progress_token) + reset_current_client_id(client_id_token) if not torch.is_tensor(images) or images.ndim < 1 or images.shape[0] != len(group): raise RuntimeError("Continuous VAE decoder returned an invalid batch shape") return images diff --git a/execution.py b/execution.py index 5ced90944..8d76b3019 100644 --- a/execution.py +++ b/execution.py @@ -437,7 +437,7 @@ def _send_cached_ui(server, client_id, node_id, display_node_id, cached, prompt_ cached_ui = cached.ui or {} server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, client_id) -async def execute(server, client_id, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs): +async def execute(server, client_id, legacy_server_state, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs): unique_id = current_item real_node_id = dynprompt.get_real_node_id(unique_id) display_node_id = dynprompt.get_display_node_id(unique_id) @@ -497,7 +497,7 @@ async def execute(server, client_id, dynprompt, caches, current_item, extra_data get_progress_state().start_progress(unique_id) input_data_all, missing_keys, v3_data = get_input_data(inputs, class_def, unique_id, execution_list, dynprompt, extra_data) if client_id is not None: - if not getattr(getattr(server, "prompt_queue", None), "cooperative", False): + if legacy_server_state: server.last_node_id = display_node_id server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id }, client_id) @@ -667,11 +667,12 @@ async def execute(server, client_id, dynprompt, caches, current_item, extra_data return (ExecutionResult.SUCCESS, None, None) class PromptExecutor: - def __init__(self, server, cache_type=False, cache_args=None, shared_outputs=None): + def __init__(self, server, cache_type=False, cache_args=None, shared_outputs=None, legacy_server_state=False): self.cache_args = cache_args self.cache_type = cache_type self.server = server self.shared_outputs = shared_outputs + self.legacy_server_state = legacy_server_state self.reset() def reset(self): @@ -739,9 +740,10 @@ class PromptExecutor: cooperative = getattr(getattr(self.server, "prompt_queue", None), "cooperative", False) if not cooperative: nodes.interrupt_processing(False) + legacy_server_state = not cooperative or self.legacy_server_state self.client_id = extra_data.get("client_id") - if not cooperative: + if legacy_server_state: self.server.client_id = self.client_id client_id_token = set_current_client_id(self.client_id) @@ -795,7 +797,7 @@ class PromptExecutor: break assert node_id is not None, "Node ID should not be None at this point" - result, error, ex = await execute(self.server, self.client_id, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs) + result, error, ex = await execute(self.server, self.client_id, legacy_server_state, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs) self.success = result != ExecutionResult.FAILURE if result == ExecutionResult.FAILURE: self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex) @@ -841,13 +843,15 @@ class PromptExecutor: "outputs": ui_outputs, "meta": meta_outputs, } - if not cooperative: + if legacy_server_state: self.server.last_node_id = None if comfy.model_management.DISABLE_SMART_MEMORY: comfy.model_management.unload_all_models() finally: for cache in self.caches.all: cache.release_prompt() + if self.legacy_server_state: + self.server.last_node_id = None if not cooperative: comfy.memory_management.set_ram_cache_release_state(None, 0) self._notify_prompt_lifecycle("end", prompt_id) diff --git a/main.py b/main.py index cf1bc2a0d..39bf56e58 100644 --- a/main.py +++ b/main.py @@ -341,8 +341,14 @@ def prompt_executor_config(): return cache_type, {"lru": args.cache_lru, "ram": cache_ram, "ram_inactive": cache_ram_inactive} -async def execute_prompt_async(q, server_instance, item, item_id, cache_type, cache_args, shared_outputs): - executor = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args=cache_args, shared_outputs=shared_outputs) +async def execute_prompt_async(q, server_instance, item, item_id, cache_type, cache_args, shared_outputs, legacy_server_state=False): + executor = execution.PromptExecutor( + server_instance, + cache_type=cache_type, + cache_args=cache_args, + shared_outputs=shared_outputs, + legacy_server_state=legacy_server_state, + ) execution_start_time = time.perf_counter() prompt_id = item[1] server_instance.last_prompt_id = prompt_id @@ -423,6 +429,14 @@ def continuous_prompt_key(item): if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str) and value[0] in prompt and isinstance(value[1], (int, float)): pending.append(value[0]) + for node_id in dependencies: + node_cls = nodes.NODE_CLASS_MAPPINGS.get(prompt[node_id].get("class_type")) + if node_cls is None: + return None + module = getattr(node_cls, "RELATIVE_PYTHON_MODULE", None) + if isinstance(module, str) and (module == "custom_nodes" or module.startswith("custom_nodes.") or module == "comfy_api_nodes" or module.startswith("comfy_api_nodes.")): + return None + sampler_ids = [node_id for node_id in dependencies if prompt[node_id].get("class_type") in comfy.continuous_batching.CONTINUOUS_SAMPLER_NODE_FAMILIES] if len(sampler_ids) != 1: return None @@ -467,7 +481,7 @@ async def cooperative_prompt_worker(q, server_instance, max_prompts): active_key = item_key family = comfy.continuous_batching.CONTINUOUS_SAMPLER_NODE_FAMILIES.get(item_key[0]) if item_key is not None else None group_cache_scope = comfy.text_encoders.anima_cache.begin_cache_scope(family == comfy.continuous_batching.FAMILY_ANIMA) - active.add(asyncio.create_task(execute_prompt_async(q, server_instance, item, item_id, cache_type, cache_args, shared_outputs))) + active.add(asyncio.create_task(execute_prompt_async(q, server_instance, item, item_id, cache_type, cache_args, shared_outputs, item_key is None))) if item_key is None: break diff --git a/nodes.py b/nodes.py index f7b573081..61c85861b 100644 --- a/nodes.py +++ b/nodes.py @@ -367,6 +367,9 @@ class ContinuousVAEDecode: max_batch_size=max_batch_size, admission_delay=admission_delay_ms / 1000.0, prompt_id=execution_context.prompt_id if execution_context is not None else None, + node_id=execution_context.node_id if execution_context is not None else None, + client_id=get_current_client_id(), + progress_registry=get_progress_state(), ) images = await comfy.continuous_batching.decode_vae_continuous(state) if len(images.shape) == 5: #Combine batches diff --git a/tests-unit/comfy_test/test_continuous_batching.py b/tests-unit/comfy_test/test_continuous_batching.py index a5b65bfbe..502fb7fe7 100644 --- a/tests-unit/comfy_test/test_continuous_batching.py +++ b/tests-unit/comfy_test/test_continuous_batching.py @@ -9,6 +9,7 @@ import comfy.continuous_batching as continuous_batching import comfy.model_base import comfy.patcher_extension import comfy.supported_models +import comfy.utils from comfy.continuous_batching import ( FAMILY_ANIMA, FAMILY_SD15, @@ -28,8 +29,8 @@ from comfy.continuous_batching import ( euler_step, ) from comfy_execution.graph import DynamicPrompt -from comfy_execution.progress import ProgressRegistry, get_progress_state, reset_progress_state -from comfy_execution.utils import CurrentNodeContext, get_current_client_id, get_executing_context, reset_current_client_id, set_current_client_id +from comfy_execution.progress import ProgressRegistry, get_progress_state, reset_progress_registry, set_progress_registry, reset_progress_state +from comfy_execution.utils import get_current_client_id, get_executing_context, reset_current_client_id, set_current_client_id class FakeState: @@ -523,6 +524,32 @@ def test_failure_clears_requests_and_finishing_request_reloads_survivor(): asyncio.run(residency()) +def test_cancelled_final_sampler_request_is_interrupted_without_cancelling_peer(): + cancelled = set() + continuous_batching.set_cancel_checker(cancelled.__contains__) + + async def run(): + cancelled_state = FakeState("cancelled", 1) + cancelled_state.prompt_id = "cancelled" + peer_state = FakeState("peer", 1) + peer_state.prompt_id = "peer" + coordinator = ContinuousBatchCoordinator("key", cancelled_state) + coordinator.session = FakeSession(on_first_step=lambda: cancelled.add("cancelled")) + results = await asyncio.gather( + coordinator.submit(cancelled_state), + coordinator.submit(peer_state), + return_exceptions=True, + ) + assert isinstance(results[0], comfy.model_management.InterruptProcessingException) + assert results[1] == "peer output" + assert cancelled_state.cleared and peer_state.cleared + + try: + asyncio.run(run()) + finally: + continuous_batching.set_cancel_checker(None) + + def test_continuous_vae_single_request_passes_the_original_tensor(monkeypatch): monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) @@ -724,20 +751,46 @@ def test_continuous_vae_node_preserves_standard_five_dimensional_image_contract( asyncio.run(run()) -def test_continuous_vae_coordinator_runs_without_the_submitter_context(monkeypatch): +def test_continuous_vae_decode_uses_representative_context_and_progress_registry(monkeypatch): seen = [] + request_registry = ProgressRegistry("request", DynamicPrompt({})) + stale_registry = ProgressRegistry("stale", DynamicPrompt({})) class ContextVAE(FakeVAE): def decode(self, samples): - seen.append(get_executing_context()) + context = get_executing_context() + seen.append((context, get_current_client_id(), get_progress_state())) + comfy.utils.ProgressBar(1).update(1) return super().decode(samples) monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", ContextVAE) + monkeypatch.setattr( + comfy.utils, + "PROGRESS_BAR_HOOK", + lambda value, total, preview, node_id=None: get_progress_state().update_progress( + node_id or get_executing_context().node_id, value, total, preview + ), + ) async def run(): vae = ContextVAE() - with CurrentNodeContext("request", "decode"): - await decode_vae_continuous(_vae_request(vae, torch.zeros(1, 4, 2, 2))) + progress_token = set_progress_registry(stale_registry) + client_token = set_current_client_id("stale-client") + try: + state = _vae_request(vae, torch.zeros(1, 4, 2, 2), prompt_id="request") + state.node_id = "decode" + state.client_id = "request-client" + state.progress_registry = request_registry + await decode_vae_continuous(state) + assert state.progress_registry is None + finally: + reset_current_client_id(client_token) + reset_progress_registry(progress_token) asyncio.run(run()) - assert seen == [None] + assert seen[0][0].prompt_id == "request" + assert seen[0][0].node_id == "decode" + assert seen[0][1] == "request-client" + assert seen[0][2] is request_registry + assert request_registry.nodes["decode"]["value"] == 1 + assert stale_registry.nodes == {} diff --git a/tests-unit/execution_test/test_continuous_worker.py b/tests-unit/execution_test/test_continuous_worker.py index 6e1b7be71..6b07465ef 100644 --- a/tests-unit/execution_test/test_continuous_worker.py +++ b/tests-unit/execution_test/test_continuous_worker.py @@ -150,3 +150,104 @@ def test_continuous_prompt_key_tracks_model_graph_and_preview_method(): first_key = main.continuous_prompt_key(_queue_item(first, ["output"], "auto")) assert first_key != main.continuous_prompt_key(_queue_item(first, ["output"], "none")) assert first_key != main.continuous_prompt_key(_queue_item(second, ["output"], "auto")) + + +def test_continuous_prompt_key_rejects_custom_and_api_nodes_but_allows_core_modules(monkeypatch): + prompt = _continuous_prompt(["sampler"]) + prompt["output"]["inputs"]["images"] = ["sampler", 0] + + for module, expected_none in ( + ("custom_nodes", True), + ("custom_nodes.extension", True), + ("comfy_api_nodes", True), + ("comfy_api_nodes.extension", True), + (None, False), + ("comfy_extras.nodes_custom_sampler", False), + ): + node_cls = type("Dependency", (), {"RELATIVE_PYTHON_MODULE": module}) + monkeypatch.setitem(main.nodes.NODE_CLASS_MAPPINGS, "ContinuousKeyDependency", node_cls) + prompt["model"]["class_type"] = "ContinuousKeyDependency" + key = main.continuous_prompt_key(_queue_item(prompt, ["output"])) + assert (key is None) is expected_none + + prompt["model"]["class_type"] = "MissingContinuousKeyDependency" + assert main.continuous_prompt_key(_queue_item(prompt, ["output"])) is None + + +def test_cooperative_serial_prompt_owns_legacy_server_state(monkeypatch): + observed = [] + + class LegacyNode: + @classmethod + def INPUT_TYPES(cls): + return {"required": {}} + + RETURN_TYPES = () + FUNCTION = "run" + + def run(self): + observed.append((server.client_id, server.last_node_id)) + return () + + class Queue: + def __init__(self, item): + self.item = item + self.returned = False + self.cooperative = False + + def set_cooperative(self, enabled): + self.cooperative = enabled + + def is_cancelled(self, prompt_id): + return False + + def get_if(self, predicate): + if not self.returned: + self.returned = True + assert predicate(self.item) + return self.item, 0 + raise StopWorker() + + def task_done(self, *args, **kwargs): + pass + + def finish_cooperative_drain(self): + pass + + def get_flags(self): + return {} + + class Server: + def __init__(self, queue): + self.prompt_queue = queue + self.client_id = "stale-client" + self.last_node_id = "stale-node" + self.last_prompt_id = None + + def send_sync(self, event, data, client_id): + pass + + prompt = {"legacy": {"class_type": "LegacyServerStateProbe", "inputs": {}}} + queue = Queue((0, "legacy-prompt", prompt, {"client_id": "legacy-client"}, ["legacy"], {})) + server = Server(queue) + monkeypatch.setitem(main.nodes.NODE_CLASS_MAPPINGS, "LegacyServerStateProbe", LegacyNode) + monkeypatch.setattr(main, "prompt_executor_config", lambda: (execution.CacheType.NONE, {"lru": 0, "ram": 0, "ram_inactive": 0})) + monkeypatch.setattr(main.comfy.memory_management, "set_ram_cache_release_state", lambda *args: None) + monkeypatch.setattr(main.comfy.continuous_batching, "set_cancel_checker", lambda checker: None) + monkeypatch.setattr(main.comfy.model_management, "soft_empty_cache", lambda: None) + monkeypatch.setattr(main.hook_breaker_ac10a0, "restore_functions", lambda: None) + monkeypatch.setattr(main.gc, "collect", lambda: None) + monkeypatch.setattr(main.asset_seeder, "is_disabled", lambda: True) + monkeypatch.setattr(main.asset_seeder, "pause", lambda: None) + monkeypatch.setattr(main.asset_seeder, "resume", lambda: None) + + async def run(): + try: + await main.cooperative_prompt_worker(queue, server, 2) + except StopWorker: + pass + + asyncio.run(run()) + assert observed == [("legacy-client", "legacy")] + assert server.client_id == "legacy-client" + assert server.last_node_id is None