diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index 9145c688da85..f36a201000c4 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -79,7 +79,9 @@ class Lineariser: Linearizer that ensures ordered delivery from multiple concurrent producers. Creates one input channel per producer and streams messages to output - in sequence-number order, buffering only out-of-order arrivals. + in sequence-number order. Each producer must provide a monotonic + increasing order of sequence numbers. For best performance, sequence + numbers should be assigned round-robin to producers. """ def __init__( @@ -89,6 +91,18 @@ def __init__( self.ch_out = ch_out self.num_producers = num_producers self.input_channels = [context.create_channel() for _ in range(num_producers)] + self._producer_slots = [asyncio.Semaphore(1) for _ in range(num_producers)] + + async def acquire(self, producer_id: int) -> Channel[TableChunk]: + """ + Wait for capacity to produce, then return the producer's channel. + + Capacity is returned only after the lineariser has forwarded the + producer's message downstream. Acquiring before constructing the next + message therefore bounds each producer to one in-flight message. + """ + await self._producer_slots[producer_id].acquire() + return self.input_channels[producer_id] async def drain(self) -> None: """ @@ -101,7 +115,8 @@ async def drain(self) -> None: buffer = {} pending_tasks = { - asyncio.create_task(ch.recv(self.context)): ch for ch in self.input_channels + asyncio.create_task(ch.recv(self.context)): producer_id + for producer_id, ch in enumerate(self.input_channels) } while pending_tasks: @@ -110,22 +125,27 @@ async def drain(self) -> None: ) for task in done: - ch = pending_tasks.pop(task) + producer_id = pending_tasks.pop(task) msg = await task if msg is not None: - buffer[msg.sequence_number] = msg - new_task = asyncio.create_task(ch.recv(self.context)) - pending_tasks[new_task] = ch + buffer[msg.sequence_number] = (msg, producer_id) # Forward consecutive messages while next_seq in buffer: - await self.ch_out.send(self.context, buffer.pop(next_seq)) + msg, producer_id = buffer.pop(next_seq) + await self.ch_out.send(self.context, msg) + self._producer_slots[producer_id].release() + ch = self.input_channels[producer_id] + new_task = asyncio.create_task(ch.recv(self.context)) + pending_tasks[new_task] = producer_id next_seq += 1 # Forward any remaining buffered messages for seq in sorted(buffer.keys()): - await self.ch_out.send(self.context, buffer.pop(seq)) + msg, producer_id = buffer.pop(seq) + await self.ch_out.send(self.context, msg) + self._producer_slots[producer_id].release() await self.ch_out.drain(self.context) @@ -264,8 +284,9 @@ async def dataframescan_node( producer_id = task_idx % num_producers producer_tasks[producer_id].append((task_idx, ir_slice)) - async def _producer(producer_id: int, ch_out: Channel) -> None: + async def _producer(producer_id: int) -> None: for task_idx, ir_slice in producer_tasks[producer_id]: + ch_out = await lineariser.acquire(producer_id) await read_chunk( context, ir_slice, @@ -282,10 +303,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ): await gather_in_task_group( lineariser.drain(), - *( - _producer(i, ch_in) - for i, ch_in in enumerate(lineariser.input_channels) - ), + *(_producer(i) for i in range(num_producers)), ) @@ -659,8 +677,9 @@ async def scan_node( # mypy resolves __iter__ on union-of-sequences to the common base (IR) producer_tasks[producer_id].append((task_idx, scan)) # type: ignore[arg-type] - async def _producer(producer_id: int, ch_out: Channel) -> None: + async def _producer(producer_id: int) -> None: for task_idx, scan in producer_tasks[producer_id]: + ch_out = await lineariser.acquire(producer_id) await read_chunk( context, scan, @@ -677,10 +696,7 @@ async def _producer(producer_id: int, ch_out: Channel) -> None: ): await gather_in_task_group( lineariser.drain(), - *( - _producer(i, ch_in) - for i, ch_in in enumerate(lineariser.input_channels) - ), + *(_producer(i) for i in range(num_producers)), ) diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index a14ff016e985..6eb99da2b0bd 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -17,8 +17,11 @@ import polars as pl from cudf_streaming.table_chunk import TableChunk +from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk +from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame +from cudf_polars.streaming.actor_graph.io import Lineariser from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk if TYPE_CHECKING: @@ -68,6 +71,55 @@ async def send_and_recv(): assert tracer.row_count == 3 +@pytest.mark.spmd +def test_lineariser_backpressures_each_producer(spmd_engine: SPMDEngine) -> None: + context = spmd_engine.context + ch_out = context.create_channel() + lineariser = Lineariser(context, ch_out, num_producers=2) + produced: list[list[int]] = [[], []] + output: list[int] = [] + + async def run() -> list[list[int]]: + release_gap = asyncio.Event() + out_of_order_sent = asyncio.Event() + + async def producer(producer_id: int, sequence_numbers: list[int]) -> None: + if producer_id == 1: + await release_gap.wait() + for sequence_number in sequence_numbers: + ch_in = await lineariser.acquire(producer_id) + produced[producer_id].append(sequence_number) + await ch_in.send( + context, + Message(sequence_number, ArbitraryChunk(sequence_number)), + ) + if sequence_number == 2: + out_of_order_sent.set() + await lineariser.input_channels[producer_id].drain(context) + + async def consumer() -> None: + while (msg := await ch_out.recv(context)) is not None: + output.append(ArbitraryChunk.from_message(msg).release()) + + async with asyncio.TaskGroup() as tg: + tg.create_task(lineariser.drain()) + tg.create_task(producer(0, [0, 2, 4])) + tg.create_task(producer(1, [1, 3, 5])) + tg.create_task(consumer()) + + await out_of_order_sent.wait() + await asyncio.sleep(0) + produced_before_gap = [values.copy() for values in produced] + release_gap.set() + + return produced_before_gap + + produced_before_gap = asyncio.run(run()) + + assert produced_before_gap == [[0, 2], []] + assert output == list(range(6)) + + def test_structlog_streaming_node_events(timeout_seconds: int): """Test that structlog emits 'Streaming Actor' events when tracing is enabled.""" pytest.importorskip("structlog")