ComfyUI/tests/execution/testing_nodes/testing-pack/async_test_nodes.py

import torch
import asyncio
from typing import Dict
from comfy.utils import ProgressBar
from comfy_execution.graph_utils import GraphBuilder
from comfy.comfy_types.node_typing import ComfyNodeABC
from comfy.comfy_types import IO


class TestAsyncValidation(ComfyNodeABC):
    """Test node with async VALIDATE_INPUTS."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": ("FLOAT", {"default": 5.0}),
                "threshold": ("FLOAT", {"default": 10.0}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process"
    CATEGORY = "experimental/async"

    @classmethod
    async def VALIDATE_INPUTS(cls, value, threshold):
        # Simulate async validation (e.g., checking remote service)
        await asyncio.sleep(0.05)

        if value > threshold:
            return f"Value {value} exceeds threshold {threshold}"
        return True

    def process(self, value, threshold):
        # Create image based on value
        intensity = value / 10.0
        image = torch.ones([1, 512, 512, 3]) * intensity
        return (image,)


class TestAsyncError(ComfyNodeABC):
    """Test node that errors during async execution."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": (IO.ANY, {}),
                "error_after": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 10.0}),
            },
        }

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "error_execution"
    CATEGORY = "experimental/async"

    async def error_execution(self, value, error_after):
        await asyncio.sleep(error_after)
        raise RuntimeError("Intentional async execution error for testing")


class TestAsyncValidationError(ComfyNodeABC):
    """Test node with async validation that always fails."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": ("FLOAT", {"default": 5.0}),
                "max_value": ("FLOAT", {"default": 10.0}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process"
    CATEGORY = "experimental/async"

    @classmethod
    async def VALIDATE_INPUTS(cls, value, max_value):
        await asyncio.sleep(0.05)
        # Always fail validation for values > max_value
        if value > max_value:
            return f"Async validation failed: {value} > {max_value}"
        return True

    def process(self, value, max_value):
        # This won't be reached if validation fails
        image = torch.ones([1, 512, 512, 3]) * (value / max_value)
        return (image,)


class TestAsyncTimeout(ComfyNodeABC):
    """Test node that simulates timeout scenarios."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": (IO.ANY, {}),
                "timeout": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0}),
                "operation_time": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 10.0}),
            },
        }

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "timeout_execution"
    CATEGORY = "experimental/async"

    async def timeout_execution(self, value, timeout, operation_time):
        try:
            # This will timeout if operation_time > timeout
            await asyncio.wait_for(asyncio.sleep(operation_time), timeout=timeout)
            return (value,)
        except asyncio.TimeoutError:
            raise RuntimeError(f"Operation timed out after {timeout} seconds")


class TestSyncError(ComfyNodeABC):
    """Test node that errors synchronously (for mixed sync/async testing)."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": (IO.ANY, {}),
            },
        }

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "sync_error"
    CATEGORY = "experimental/async"

    def sync_error(self, value):
        raise RuntimeError("Intentional sync execution error for testing")


class TestOOMError(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"value": (IO.ANY, {})}}

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "oom_error"
    CATEGORY = "experimental/async"

    def oom_error(self, value):
        raise torch.OutOfMemoryError("Intentional out of memory error for testing")


class TestMixedExpansionFailure(ComfyNodeABC):
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"value": ("INT", {})}}

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "expand"
    CATEGORY = "experimental/async"

    def expand(self, value):
        image = torch.zeros([1, 32, 32, 3])
        if value == 0:
            return (image,)

        graph = GraphBuilder()
        error = graph.node("TestSyncError", value=image)
        return {
            "result": (error.out(0),),
            "expand": graph.finalize(),
        }


class TestMalformedExpansion(ComfyNodeABC):
    """Expands to a graph referencing a missing node class, so the failure
    happens in the executor after the node function has returned."""

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"value": (IO.ANY, {})}}

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "expand"
    CATEGORY = "experimental/async"

    def expand(self, value):
        graph = GraphBuilder()
        missing = graph.node("TestNodeClassThatDoesNotExist", value=value)
        return {
            "result": (missing.out(0),),
            "expand": graph.finalize(),
        }


class TestMalformedResult(ComfyNodeABC):
    """Returns a non-tuple result so the failure happens while the executor
    merges results, after the node function has returned."""

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"value": (IO.ANY, {})}}

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "run"
    CATEGORY = "experimental/async"

    def run(self, value):
        return 5


class TestCyclicExpansion(ComfyNodeABC):
    """Expands to an output node that consumes this node's own pending output,
    forming a cycle through the expansion completion link."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {"value": (IO.ANY, {})},
            "hidden": {"unique_id": "UNIQUE_ID"},
        }

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "expand"
    CATEGORY = "experimental/async"

    def expand(self, value, unique_id):
        graph = GraphBuilder()
        graph.node("TestAsyncOutput", value=[unique_id, 0], seconds=0.0)
        return {
            "result": (value,),
            "expand": graph.finalize(),
        }


class TestExpansionWithFailingOutput(ComfyNodeABC):
    """Expands to a subgraph whose result succeeds while a side branch ending
    in an output node fails."""

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"image": (IO.IMAGE, {})}}

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "expand"
    CATEGORY = "experimental/async"

    def expand(self, image):
        graph = GraphBuilder()
        error = graph.node("TestSyncError", value=image)
        graph.node("PreviewImage", images=error.out(0))
        passthrough = graph.node("StubImage", content="WHITE", height=32, width=32, batch_size=1)
        return {
            "result": (passthrough.out(0),),
            "expand": graph.finalize(),
        }


class TestAsyncOutput(ComfyNodeABC):
    """Async output node with no return sockets, used to test partial failure
    handling across pending async invocations."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": (IO.ANY, {}),
                "seconds": (IO.FLOAT, {"default": 0.1}),
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "run"
    CATEGORY = "experimental/async"

    async def run(self, value, seconds=0.1):
        await asyncio.sleep(seconds)
        return {"ui": {"values": [1]}}


class TestAsyncLazyCheck(ComfyNodeABC):
    """Test node with async check_lazy_status."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "input1": (IO.ANY, {"lazy": True}),
                "input2": (IO.ANY, {"lazy": True}),
                "condition": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process"
    CATEGORY = "experimental/async"

    async def check_lazy_status(self, condition, input1, input2):
        # Simulate async checking (e.g., querying remote service)
        await asyncio.sleep(0.05)

        needed = []
        if condition and input1 is None:
            needed.append("input1")
        if not condition and input2 is None:
            needed.append("input2")
        return needed

    def process(self, input1, input2, condition):
        # Return a simple image
        return (torch.ones([1, 512, 512, 3]),)


class TestDynamicAsyncGeneration(ComfyNodeABC):
    """Test node that dynamically generates async nodes."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image1": ("IMAGE",),
                "image2": ("IMAGE",),
                "num_async_nodes": ("INT", {"default": 3, "min": 1, "max": 10}),
                "sleep_duration": ("FLOAT", {"default": 0.2, "min": 0.1, "max": 1.0}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate_async_workflow"
    CATEGORY = "experimental/async"

    def generate_async_workflow(self, image1, image2, num_async_nodes, sleep_duration):
        g = GraphBuilder()

        # Create multiple async sleep nodes
        sleep_nodes = []
        for i in range(num_async_nodes):
            image = image1 if i % 2 == 0 else image2
            sleep_node = g.node("TestSleep", value=image, seconds=sleep_duration)
            sleep_nodes.append(sleep_node)

        # Average all results
        if len(sleep_nodes) == 1:
            final_node = sleep_nodes[0]
        else:
            avg_inputs = {"input1": sleep_nodes[0].out(0)}
            for i, node in enumerate(sleep_nodes[1:], 2):
                avg_inputs[f"input{i}"] = node.out(0)
            final_node = g.node("TestVariadicAverage", **avg_inputs)

        return {
            "result": (final_node.out(0),),
            "expand": g.finalize(),
        }


class TestAsyncResourceUser(ComfyNodeABC):
    """Test node that uses resources during async execution."""

    # Class-level resource tracking for testing
    _active_resources: Dict[str, bool] = {}

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": (IO.ANY, {}),
                "resource_id": ("STRING", {"default": "resource_0"}),
                "duration": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0}),
            },
        }

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "use_resource"
    CATEGORY = "experimental/async"

    async def use_resource(self, value, resource_id, duration):
        # Check if resource is already in use
        if self._active_resources.get(resource_id, False):
            raise RuntimeError(f"Resource {resource_id} is already in use!")

        # Mark resource as in use
        self._active_resources[resource_id] = True

        try:
            # Simulate resource usage
            await asyncio.sleep(duration)
            return (value,)
        finally:
            # Always clean up resource
            self._active_resources[resource_id] = False


class TestAsyncBatchProcessing(ComfyNodeABC):
    """Test async processing of batched inputs."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "process_time_per_item": ("FLOAT", {"default": 0.1, "min": 0.01, "max": 1.0}),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "process_batch"
    CATEGORY = "experimental/async"

    async def process_batch(self, images, process_time_per_item, unique_id):
        batch_size = images.shape[0]
        pbar = ProgressBar(batch_size, node_id=unique_id)

        # Process each image in the batch
        processed = []
        for i in range(batch_size):
            # Simulate async processing
            await asyncio.sleep(process_time_per_item)

            # Simple processing: invert the image
            processed_image = 1.0 - images[i:i+1]
            processed.append(processed_image)

            pbar.update(1)

        # Stack processed images
        result = torch.cat(processed, dim=0)
        return (result,)


class TestAsyncConcurrentLimit(ComfyNodeABC):
    """Test concurrent execution limits for async nodes."""

    _semaphore = asyncio.Semaphore(2)  # Only allow 2 concurrent executions

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "value": (IO.ANY, {}),
                "duration": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 2.0}),
                "node_id": ("INT", {"default": 0}),
            },
        }

    RETURN_TYPES = (IO.ANY,)
    FUNCTION = "limited_execution"
    CATEGORY = "experimental/async"

    async def limited_execution(self, value, duration, node_id):
        async with self._semaphore:
            # Node {node_id} acquired semaphore
            await asyncio.sleep(duration)
            # Node {node_id} releasing semaphore
            return (value,)


# Add node mappings
ASYNC_TEST_NODE_CLASS_MAPPINGS = {
    "TestAsyncValidation": TestAsyncValidation,
    "TestAsyncError": TestAsyncError,
    "TestAsyncValidationError": TestAsyncValidationError,
    "TestAsyncTimeout": TestAsyncTimeout,
    "TestSyncError": TestSyncError,
    "TestOOMError": TestOOMError,
    "TestMixedExpansionFailure": TestMixedExpansionFailure,
    "TestMalformedExpansion": TestMalformedExpansion,
    "TestMalformedResult": TestMalformedResult,
    "TestCyclicExpansion": TestCyclicExpansion,
    "TestExpansionWithFailingOutput": TestExpansionWithFailingOutput,
    "TestAsyncOutput": TestAsyncOutput,
    "TestAsyncLazyCheck": TestAsyncLazyCheck,
    "TestDynamicAsyncGeneration": TestDynamicAsyncGeneration,
    "TestAsyncResourceUser": TestAsyncResourceUser,
    "TestAsyncBatchProcessing": TestAsyncBatchProcessing,
    "TestAsyncConcurrentLimit": TestAsyncConcurrentLimit,
}

ASYNC_TEST_NODE_DISPLAY_NAME_MAPPINGS = {
    "TestAsyncValidation": "Test Async Validation",
    "TestAsyncError": "Test Async Error",
    "TestAsyncValidationError": "Test Async Validation Error",
    "TestAsyncTimeout": "Test Async Timeout",
    "TestSyncError": "Test Sync Error",
    "TestOOMError": "Test OOM Error",
    "TestMixedExpansionFailure": "Test Mixed Expansion Failure",
    "TestMalformedExpansion": "Test Malformed Expansion",
    "TestMalformedResult": "Test Malformed Result",
    "TestCyclicExpansion": "Test Cyclic Expansion",
    "TestExpansionWithFailingOutput": "Test Expansion With Failing Output",
    "TestAsyncOutput": "Test Async Output",
    "TestAsyncLazyCheck": "Test Async Lazy Check",
    "TestDynamicAsyncGeneration": "Test Dynamic Async Generation",
    "TestAsyncResourceUser": "Test Async Resource User",
    "TestAsyncBatchProcessing": "Test Async Batch Processing",
    "TestAsyncConcurrentLimit": "Test Async Concurrent Limit",
}