From c3f1e92e4e2113bb39a8d190d6d6c1332f0d6348 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 8 Sep 2026 15:27:28 +0100 Subject: [PATCH] Ensure that Lineariser doesn't buffer unboundedly The previous approach to buffered refill in the Lineariser could result in unbounded buffering of tasks. If the next sequence number tasks was slow, we would still arbitrarily refill from fast producers, resulting in excessive memory pressure and removing the backpressure the finite capacity channels are intended to provide. To fix this, give each producer a size-1 semaphore that it must acquire before being allowed to kick off a memory-using task. This semaphore is only released once the matching message has been forwarded into the downstream channel. This doesn't change the reordering properties of the lineariser, but does make it more important to assign message sequence ids in round-robin fashion to the producer tasks to ensure performance is good. --- .../cudf_polars/streaming/actor_graph/io.py | 52 ++++++++++++------- .../tests/streaming/test_tracing.py | 52 +++++++++++++++++++ 2 files changed, 86 insertions(+), 18 deletions(-) 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")