Fix continuous batching request isolation

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

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)

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
 

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

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 == {}

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