Flexible Edge Behaviors Between Nodes

Whereas each node in the graph is defined by an asynchronous generator function, the edge between two nodes is represented by a queue object that feeds data into the destination node and collects data from the source nodes (see Tasks and Queues). By default, this queue object spits out data items to the associated destination node one at a time, as it becomes available from any of the source nodes, and as efficieintly as coordinated by the async event loop. In your application, you may want to override this default queue for custom edge behaviors for one or more edges in the graph.

When a node is created with add_node(), custom edge behaviors can be implemented by passing in a custom queue object to the queue parameter. This queue object must be an instance of asyncio.Queue or a subclass thereof. (The default queue object is simply an instance of asyncio.Queue.) A custom queue object can result in a variety of edge behaviors between nodes:

  • A queue pre-loaded with data items:

    At the outset of the graph execution, if you would like certain nodes (including non-starting nodes of the graph) to receive specific data items immediately, you can pass in a queue object pre-loaded with these data items for these nodes. Implementationally, create an asyncio.Queue instance and call await queue.put(item) as desired to pre-load the queue with data items, then pass this queue object to the queue parameter of the destination node at add_node().

  • Do something with the data after it is received from a source node and before it is fed to the destination node:

    If you would like to transform the data or perform any other operation on it, your custom queue likely comes from a subclass of asyncio.Queue where you override the get and/or put methods, becase these methods control what happens to the data item after it has arrived at the queue and before it is fed to the destination node.

In the following examples, we focus on the second use case, where flexible edge behaviors are achieved by a subclass of asyncio.Queue.

Batching

Instead of the edge queue receiving and feeding data items one at a time, you may want to batch the data items before feeding them to the destination node. To do so, let’s define the following:

  • A BatchQueue class that subclasses asyncio.Queue for batching data items.

  • A sentinel class EndOfData to signal the end of data.

Perhaps the most common use case for batching is to group data items into batches of a fixed size:

        flowchart LR

start1[ ] -.- queue1((" "))
subgraph data_source
    queue1 --> node1["  "]
end
style start1 fill-opacity:0, stroke-opacity:0;

queue2((custom queue<br/>to batch data))
node2["async gen<br/>function"]
subgraph batched_inputs
    queue2 --> |"three inputs:<br/>[1, 2, 3, 4]<br/>[5, 6, 7, 8]<br/>[9, 10]"| node2
end
style queue2 fill:#ffccff, stroke:#030303, stroke-width:2px;

node1 --> |yields<br/>1, 2, 3, ..., 10| queue2
node2 -.-> STOP[ ]
style STOP fill-opacity:0, stroke-opacity:0;
    
import asyncio

from async_graph_data_flow import AsyncExecutor, AsyncGraph


class EndOfData:
    pass


class BatchQueue(asyncio.Queue):
    def __init__(self, batch_size: int | float):
        super().__init__()
        self.batch_size = batch_size
        self.batch: list = []

    async def put(self, item):
        # This custom queue class assumes the use of an EndOfData marker
        # to indicate the end of the data stream.
        if isinstance(item, EndOfData):
            # No more data is expected, so flush the last batch
            # if it's not empty and regardless of its size.
            if self.batch:
                await super().put(self.batch)
            # Don't forget to pass along the EndOfData marker.
            await super().put(item)
        else:
            self.batch.append(item)
            if len(self.batch) >= self.batch_size:
                await super().put(self.batch)
                self.batch = []


async def data_source():
    for i in range(1, 11):
        print(f"data_source yielding {i}")
        yield i
    # Indicate the end of the data stream,
    # so that the batch queue can flush the last batch.
    yield EndOfData()


async def batched_inputs(data):
    print(f"batched_inputs received: {data}")
    yield


if __name__ == "__main__":
    graph = AsyncGraph()

    graph.add_node(data_source)
    graph.add_node(batched_inputs, queue=BatchQueue(batch_size=4))

    graph.add_edge(data_source, batched_inputs)

    AsyncExecutor(graph).execute()

    # Output:
    # -------
    # data_source yielding 1
    # data_source yielding 2
    # data_source yielding 3
    # data_source yielding 4
    # data_source yielding 5
    # data_source yielding 6
    # data_source yielding 7
    # data_source yielding 8
    # data_source yielding 9
    # data_source yielding 10
    # batched_inputs received: [1, 2, 3, 4]
    # batched_inputs received: [5, 6, 7, 8]
    # batched_inputs received: [9, 10]
    # batched_inputs received: <__main__.EndOfData object at 0x10a990260>

Using the same BatchQueue and EndOfData defined above, it’s also possible to have the effect of waiting for all data items from the source nodes before feeding them to the destination node, by setting the batch size to float('inf') (infinity, for no batch size limit).

Beyond batch size, batching can be controlled by other criteria using your own custom queue class, such as special data items or markers (in a way, the EndOfData marker above is an example), time intervals, or other conditions.

Combining Data from Multiple Source Nodes

By default, a source node must yield data that matches the function signature of the destination node. To work around this constraint, a custom queue object at the edge between the source and destination nodes can be used. One use case is to combine data from multiple source nodes before feeding it to the destination node. For example:

        flowchart LR

start1[ ] -.- queue1(("&nbsp;"))
subgraph threes
    queue1 --> node1["&nbsp;&nbsp;"]
end
style start1 fill-opacity:0, stroke-opacity:0;

start2[ ] -.- queue2(("&nbsp;"))
subgraph fours
    queue2 --> node2["&nbsp;&nbsp;"]
end
style start2 fill-opacity:0, stroke-opacity:0;

start3[ ] -.- queue3(("&nbsp;"))
subgraph fives
    queue3 --> node3["&nbsp;&nbsp;"]
end
style start3 fill-opacity:0, stroke-opacity:0;

queue4((custom queue<br/>to combine data))
node4["async gen<br/>function"]
subgraph final_node
    queue4 --> |"three inputs<br/>to the func:<br/>(3, 4, 5)<br/>(3, 4, 5)<br/>(3, 4, 5)"| node4
end
style queue4 fill:#ffccff, stroke:#030303, stroke-width:2px;

node1 --> |yields<br/>3, 3, 3| queue4
node2 --> |yields<br/>4, 4, 4, 4| queue4
node3 --> |yields<br/>5, 5, 5, 5, 5| queue4
node4 -.-> STOP[ ]
style STOP fill-opacity:0, stroke-opacity:0;
    
import asyncio
import itertools

from async_graph_data_flow import AsyncExecutor, AsyncGraph


class CombineDataQueue(asyncio.Queue):
    def __init__(self, node_order: list[str]):
        super().__init__()
        # This custom queue class assumes that a source node yields a tuple
        # where the first element is the node marker and the rest are the data.
        # `node_order` has the node markers for all the source nodes,
        # in the order that the data from all source nodes will be combined.
        self.node_order = node_order
        self.queues = {node: asyncio.Queue() for node in node_order}

    async def put(self, item):
        node_marker, *data = item
        if node_marker not in self.node_order:
            raise ValueError(f"Node {node_marker} not in node order {self.node_order}")
        if not data:
            raise ValueError(f"Data missing for node {node_marker}")
        await self.queues[node_marker].put(data)
        # In this specific example, if and only if the individual queues
        # for *all* source nodes have data,
        # combine the data and put the result in the main queue.
        # This means that the k-th item in the main queue will be a tuple
        # of the k-th items from each source node, and that any additional data
        # from a source node will either wait until all source nodes have data
        # to create a new item for the main queue, or be ignored altogether.
        if all(not q.empty() for q in self.queues.values()):
            new = tuple(
                itertools.chain(*((q.get_nowait()) for q in self.queues.values()))
            )
            await super().put(new)


async def threes():
    for _ in range(3):
        await asyncio.sleep(0.001)
        print("threes yielding 3")
        # Yielding a tuple with the node marker and the data.
        # The node marker tells the custom queue which source node the data is from.
        yield "threes", 3


async def fours():
    for _ in range(4):
        await asyncio.sleep(0.001)
        print("fours yielding 4")
        yield "fours", 4


async def fives():
    for _ in range(5):
        await asyncio.sleep(0.001)
        print("fives yielding 5")
        yield "fives", 5


async def final_node(int1, int2, int3):
    print(f"Final node received {int1}, {int2}, {int3}")
    yield


if __name__ == "__main__":
    graph = AsyncGraph()

    graph.add_node(threes)
    graph.add_node(fours)
    graph.add_node(fives)
    graph.add_node(final_node, queue=CombineDataQueue(["threes", "fours", "fives"]))

    graph.add_edge("threes", "final_node")
    graph.add_edge("fours", "final_node")
    graph.add_edge("fives", "final_node")

    AsyncExecutor(graph).execute()

    # Output:
    # -------
    # threes yielding 3
    # fours yielding 4
    # fives yielding 5
    # Final node received 3, 4, 5
    # threes yielding 3
    # fours yielding 4
    # fives yielding 5
    # Final node received 3, 4, 5
    # threes yielding 3
    # fours yielding 4
    # fives yielding 5
    # Final node received 3, 4, 5
    # fours yielding 4
    # fives yielding 5
    # fives yielding 5