Ensure that Lineariser doesn't buffer unboundedly

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)),
             )
 
 

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")