From 2d93cac6a54c41b65daa6ef750d3e433f812cb9a Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 21 May 2026 12:38:42 -0700 Subject: [PATCH 01/26] add adjust_orderscheme utility --- .../actor_graph/collectives/orderscheme.py | 290 ++++++++++++++++ .../streaming/test_adjust_orderscheme.py | 313 ++++++++++++++++++ 2 files changed, 603 insertions(+) create mode 100644 python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py create mode 100644 python/cudf_polars/tests/streaming/test_adjust_orderscheme.py diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py new file mode 100644 index 000000000000..b849c71467c5 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -0,0 +1,290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""OrderScheme adjustment utilities for the RapidsMPF streaming runtime.""" + +from __future__ import annotations + +from collections import defaultdict +from typing import TYPE_CHECKING + +from rapidsmpf.integrations.cudf.partition import unpack_and_concat +from rapidsmpf.memory.packed_data import PackedData +from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall +from rapidsmpf.streaming.core.message import Message +from rapidsmpf.streaming.cudf.table_chunk import TableChunk + +import polars as pl + +import pylibcudf as plc +from pylibcudf.contiguous_split import pack + +from cudf_polars.containers import DataFrame, DataType +from cudf_polars.streaming.actor_graph.utils import concat_batch +from cudf_polars.utils.cuda_stream import stream_ordered_after + +if TYPE_CHECKING: + from collections.abc import Iterable + + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.memory.buffer_resource import BufferResource + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + from rapidsmpf.streaming.cudf.channel_metadata import OrderScheme + + from rmm.pylibrmm.stream import Stream + + from cudf_polars.dsl.ir import IR, IRExecutionContext + + +_PID_DTYPE = DataType(pl.Int32()) +_PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32) + + +def _contiguous_owner(pid: int, nranks: int, npartitions: int) -> int: + """Return the rank owning *pid* under contiguous partition assignment.""" + return pid * nranks // npartitions + + +def _local_partitions(rank: int, nranks: int, npartitions: int) -> list[int]: + """Return partition IDs owned by *rank* under contiguous assignment.""" + return [ + pid + for pid in range(npartitions) + if _contiguous_owner(pid, nranks, npartitions) == rank + ] + + +def _validate_schemes(input_scheme: OrderScheme, output_scheme: OrderScheme) -> None: + """Validate the first-pass flat OrderScheme adjustment contract.""" + if not output_scheme.strict_boundaries: + raise ValueError("adjust_orderscheme requires a strict output OrderScheme.") + prefix_len = len(output_scheme.keys) + if input_scheme.keys[:prefix_len] != output_scheme.keys: + raise NotImplementedError( + "adjust_orderscheme currently requires the output OrderScheme keys " + "to be a prefix of the input OrderScheme keys." + ) + + +def _split_points( + table: plc.Table, + boundary_table: plc.Table, + scheme: OrderScheme, + stream: Stream, +) -> list[int]: + """Return row split points that partition *table* by *scheme* boundaries.""" + if boundary_table.num_rows() == 0: + return [] + key_table = plc.Table([table.columns()[key.column_index] for key in scheme.keys]) + split_col = plc.search.lower_bound( + key_table, + boundary_table, + [key.order for key in scheme.keys], + [key.null_order for key in scheme.keys], + stream=stream, + ) + return ( + DataFrame.from_table( + plc.Table([split_col]), + ["split"], + [_PID_DTYPE], + stream, + ) + .to_polars()["split"] + .to_list() + ) + + +def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Table: + """Append a hidden target-partition-id column to *table*.""" + pid_col = plc.Column.from_scalar( + plc.Scalar.from_py(pid, _PID_PLC_DTYPE, stream=stream), + table.num_rows(), + stream=stream, + ) + return plc.Table([*table.columns(), pid_col]) + + +def _pack_table(table: plc.Table, stream: Stream, br: BufferResource) -> PackedData: + """Pack a pylibcudf table as a RapidsMPF PackedData payload.""" + return PackedData.from_cudf_packed_columns( + pack(table, stream, mr=br.device_mr), + stream, + br, + ) + + +def _unpack_remote_piece( + packed: PackedData, + stream: Stream, + br: BufferResource, +) -> tuple[int, TableChunk] | None: + """Unpack one remote piece and recover its hidden target partition ID.""" + table = unpack_and_concat([packed], stream=stream, br=br) + if table.num_rows() == 0: + return None + *payload_cols, pid_col = table.columns() + pid = int( + DataFrame.from_table( + plc.Table([pid_col]), + ["pid"], + [_PID_DTYPE], + stream, + ) + .to_polars() + .item(0, 0) + ) + payload = plc.concatenate.concatenate( + [plc.Table(payload_cols)], stream=stream, mr=br.device_mr + ) + return pid, TableChunk.from_pylibcudf_table( + payload, + stream, + exclusive_view=True, + br=br, + ) + + +def _materialize_packed_pieces( + pieces: Iterable[PackedData], + context: Context, + stream: Stream, +) -> TableChunk | None: + """Materialize packed pieces into a uniquely-owned chunk.""" + pieces = list(pieces) + if not pieces: + return None + table = unpack_and_concat(pieces, stream=stream, br=context.br()) + if table.num_rows() == 0: + return None + return TableChunk.from_pylibcudf_table( + table, + stream, + exclusive_view=True, + br=context.br(), + ) + + +async def adjust_orderscheme( + context: Context, + comm: Communicator, + ref_ir: IR, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + input_scheme: OrderScheme, + output_scheme: OrderScheme, + *, + collective_id: int | None = None, +) -> None: + """ + Adjust flat OrderScheme boundaries using contiguous partition ownership. + + Parameters + ---------- + context + The streaming context. + comm + The communicator. + ref_ir + An IR node describing the payload schema. + ir_context + The IR execution context. + ch_out + The output channel. + ch_in + The input channel. + input_scheme + The input OrderScheme. + output_scheme + The output OrderScheme. + collective_id + The collective ID to use for SparseAlltoall. + + Notes + ----- + This utility is intentionally narrow and only adjusts data messages. The + caller is responsible for receiving input metadata and sending output + metadata. + """ + _validate_schemes(input_scheme, output_scheme) + npartitions = output_scheme.num_boundaries + 1 + local_pids = _local_partitions(comm.rank, comm.nranks, npartitions) + + if comm.nranks > 1 and collective_id is None: + raise ValueError("collective_id is required when comm.nranks > 1.") + + # TODO: Narrow the peer list (avoid all-to-all control messages) + peers = [rank for rank in range(comm.nranks) if rank != comm.rank] + exchange = ( + SparseAlltoall(context, comm, collective_id, srcs=peers, dsts=peers) + if comm.nranks > 1 + else None + ) + boundary_chunk = output_scheme.get_boundaries(context.br()) + boundary_table = boundary_chunk.table_view() + local_pieces: dict[int, list[PackedData]] = defaultdict(list) + + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill( + context.br(), allow_overbooking=True + ) + if chunk.table_view().num_rows() == 0: + continue + with stream_ordered_after( + context.get_stream_from_pool, + upstreams=(chunk.stream, boundary_chunk.stream), + ) as stream: + table = chunk.table_view() + splits = _split_points(table, boundary_table, output_scheme, stream) + for pid, piece in enumerate( + plc.copying.split(table, splits, stream=stream) + ): + if piece.num_rows() == 0: + continue + owner = _contiguous_owner(pid, comm.nranks, npartitions) + if owner == comm.rank: + local_pieces[pid].append(_pack_table(piece, stream, context.br())) + else: + assert exchange is not None + exchange.insert( + owner, + _pack_table( + _append_partition_id(piece, pid, stream), + stream, + context.br(), + ), + ) + + if exchange is not None: + await exchange.insert_finished(context) + + output_chunks: dict[int, list[TableChunk]] = defaultdict(list) + for source_rank in range(comm.nranks): + if source_rank == comm.rank: + for pid, pieces in local_pieces.items(): + chunk = _materialize_packed_pieces( + pieces, context, context.get_stream_from_pool() + ) + if chunk is not None: + output_chunks[pid].append(chunk) + continue + assert exchange is not None + stream = context.get_stream_from_pool() + for packed in exchange.extract(source_rank): + remote_piece = _unpack_remote_piece(packed, stream, context.br()) + if remote_piece is None: + continue + pid, chunk = remote_piece + output_chunks[pid].append(chunk) + + for pid in local_pids: + chunks = output_chunks[pid] + chunk = ( + await concat_batch(chunks, context, ref_ir.schema, ir_context) + if chunks + else None + ) + if chunk is not None and chunk.table_view().num_rows() > 0: + await ch_out.send(context, Message(pid, chunk)) + await ch_out.drain(context) diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py new file mode 100644 index 000000000000..557b6634ebf9 --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -0,0 +1,313 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from concurrent.futures import ThreadPoolExecutor + +import pytest +from rapidsmpf.streaming.core.message import Message +from rapidsmpf.streaming.cudf.channel_metadata import ( + ChannelMetadata, + OrderKey, + OrderScheme, + Partitioning, +) +from rapidsmpf.streaming.cudf.table_chunk import TableChunk + +import polars as pl + +import pylibcudf as plc + +from cudf_polars.containers import DataFrame, DataType +from cudf_polars.dsl.ir import Empty, IRExecutionContext +from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id +from cudf_polars.streaming.actor_graph.collectives.orderscheme import ( + adjust_orderscheme, +) +from cudf_polars.streaming.actor_graph.utils import ( + gather_in_task_group, + recv_metadata, + send_metadata, +) + +_SCHEMA = {"key": DataType(pl.Int32()), "val": DataType(pl.Int32())} +_NAMES = list(_SCHEMA) +_DTYPES = list(_SCHEMA.values()) + + +def _make_scheme( + context, + boundary: int | tuple[int, ...], + *, + key_indices: tuple[int, ...] = (0,), + strict: bool = True, + stream, +) -> OrderScheme: + boundary_values = ( + boundary if isinstance(boundary, tuple) else (boundary,) * len(key_indices) + ) + boundary_df = DataFrame.from_polars( + pl.DataFrame( + { + f"k{i}": pl.Series([value], dtype=pl.Int32()) + for i, value in enumerate(boundary_values) + } + ), + stream, + ) + return OrderScheme( + [ + OrderKey(index, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE) + for index in key_indices + ], + TableChunk.from_pylibcudf_table( + boundary_df.table, + stream, + exclusive_view=False, + br=context.br(), + ), + strict_boundaries=strict, + ) + + +def _ref_ir() -> Empty: + return Empty(_SCHEMA) + + +def _local_count(comm, scheme: OrderScheme) -> int: + npartitions = scheme.num_boundaries + 1 + return sum( + pid * comm.nranks // npartitions == comm.rank for pid in range(npartitions) + ) + + +def _frame(values: list[int]) -> pl.DataFrame: + return pl.DataFrame( + { + "key": pl.Series(values, dtype=pl.Int32()), + "val": pl.Series(values, dtype=pl.Int32()), + } + ) + + +def _chunk_to_polars(chunk: TableChunk) -> pl.DataFrame: + return DataFrame.from_table( + chunk.table_view(), + _NAMES, + _DTYPES, + chunk.stream, + ).to_polars() + + +async def _adjust_and_collect( + context, + comm, + input_df: pl.DataFrame, + input_scheme: OrderScheme, + output_scheme: OrderScheme, + *, + collective_id: int | None = None, +) -> dict[int, pl.DataFrame]: + ch_in = context.create_channel() + ch_out = context.create_channel() + stream = context.get_stream_from_pool() + output: dict[int, pl.DataFrame] = {} + + async def _produce() -> None: + df = DataFrame.from_polars(input_df, stream) + await send_metadata( + ch_out, + context, + ChannelMetadata( + local_count=_local_count(comm, output_scheme), + partitioning=Partitioning(output_scheme, "inherit"), + ), + ) + await ch_in.send( + context, + Message( + comm.rank, + TableChunk.from_pylibcudf_table( + df.table, + stream, + exclusive_view=True, + br=context.br(), + ), + ), + ) + await ch_in.drain(context) + + async def _consume() -> None: + await recv_metadata(ch_out, context) + while (msg := await ch_out.recv(context)) is not None: + output[msg.sequence_number] = _chunk_to_polars( + TableChunk.from_message(msg, br=context.br()) + ) + + with ThreadPoolExecutor(max_workers=1) as executor: + ir_context = IRExecutionContext( + executor, get_cuda_stream=context.get_stream_from_pool + ) + await gather_in_task_group( + _produce(), + adjust_orderscheme( + context, + comm, + _ref_ir(), + ir_context, + ch_out, + ch_in, + input_scheme, + output_scheme, + collective_id=collective_id, + ), + _consume(), + ) + + return output + + +async def _adjust_direct( + context, + comm, + input_scheme: OrderScheme, + output_scheme: OrderScheme, + *, + collective_id: int | None = None, +) -> None: + ch_in = context.create_channel() + ch_out = context.create_channel() + with ThreadPoolExecutor(max_workers=1) as executor: + ir_context = IRExecutionContext( + executor, get_cuda_stream=context.get_stream_from_pool + ) + await adjust_orderscheme( + context, + comm, + _ref_ir(), + ir_context, + ch_out, + ch_in, + input_scheme, + output_scheme, + collective_id=collective_id, + ) + + +@pytest.mark.spmd +@pytest.mark.parametrize( + "input_keys,output_keys,strict,error,match", + [ + ((1,), (0,), True, NotImplementedError, "prefix"), + ((0,), (0,), False, ValueError, "strict output"), + ], +) +def test_adjust_orderscheme_rejects_invalid_schemes( + spmd_engine, input_keys, output_keys, strict, error, match +) -> None: + context = spmd_engine.context + stream = context.get_stream_from_pool() + input_scheme = _make_scheme(context, 4, key_indices=input_keys, stream=stream) + output_scheme = _make_scheme( + context, + 4, + key_indices=output_keys, + strict=strict, + stream=stream, + ) + + with pytest.raises(error, match=match): + asyncio.run( + _adjust_direct(context, spmd_engine.comm, input_scheme, output_scheme) + ) + + +@pytest.mark.spmd +def test_adjust_orderscheme_requires_collective_id(spmd_engine) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks == 1: + pytest.skip("collective_id is only required for multi-rank runs.") + + stream = context.get_stream_from_pool() + input_scheme = _make_scheme(context, 4, stream=stream) + output_scheme = _make_scheme(context, 4, stream=stream) + + with pytest.raises(ValueError, match="collective_id"): + asyncio.run(_adjust_direct(context, comm, input_scheme, output_scheme)) + + +@pytest.mark.spmd +@pytest.mark.parametrize( + "target_boundary,expected", + [ + (3, {0: [0, 1, 2], 1: [3, 4, 5, 6, 7]}), + (5, {0: [0, 1, 2, 3, 4], 1: [5, 6, 7]}), + ], +) +def test_adjust_orderscheme_sparse_boundary_shift( + spmd_engine, target_boundary, expected +) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks != 2: + pytest.skip("This test expects exactly two ranks.") + + keys = list(range(4)) if comm.rank == 0 else list(range(4, 8)) + stream = context.get_stream_from_pool() + # Input sorted on (key, val) is also sorted on the target key prefix. + input_scheme = _make_scheme(context, (4, 4), key_indices=(0, 1), stream=stream) + output_scheme = _make_scheme(context, target_boundary, stream=stream) + + with reserve_op_id() as op_id: + output = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(keys), + input_scheme, + output_scheme, + collective_id=op_id, + ) + ) + + result = pl.concat(output.values()) + assert result["key"].to_list() == expected[comm.rank] + assert result["val"].to_list() == expected[comm.rank] + + +@pytest.mark.spmd +@pytest.mark.parametrize( + "target_boundary,expected", + [ + (3, {0: [0, 1, 2], 1: [3, 4, 5, 6, 7]}), + (0, {1: list(range(8))}), + ], +) +def test_adjust_orderscheme_single_rank_no_collective( + spmd_engine, target_boundary, expected +) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks != 1: + pytest.skip("This test covers the single-rank path.") + + stream = context.get_stream_from_pool() + input_scheme = _make_scheme(context, 4, stream=stream) + output_scheme = _make_scheme(context, target_boundary, stream=stream) + output_by_pid = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(list(range(8))), + input_scheme, + output_scheme, + ) + ) + + assert set(output_by_pid) == set(expected) + for pid, keys in expected.items(): + assert output_by_pid[pid]["key"].to_list() == keys + assert output_by_pid[pid]["val"].to_list() == keys From cfc37d5e713348a56c1e49cc50131d6be8aa1f80 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 22 May 2026 10:14:48 -0700 Subject: [PATCH 02/26] Emit empty chunks for owned OrderScheme partitions --- .../actor_graph/collectives/orderscheme.py | 7 ++- .../streaming/test_adjust_orderscheme.py | 54 ++++++++++++++++--- 2 files changed, 51 insertions(+), 10 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py index b849c71467c5..ed971f540072 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -19,7 +19,7 @@ from pylibcudf.contiguous_split import pack from cudf_polars.containers import DataFrame, DataType -from cudf_polars.streaming.actor_graph.utils import concat_batch +from cudf_polars.streaming.actor_graph.utils import concat_batch, empty_table_chunk from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: @@ -283,8 +283,7 @@ async def adjust_orderscheme( chunk = ( await concat_batch(chunks, context, ref_ir.schema, ir_context) if chunks - else None + else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) ) - if chunk is not None and chunk.table_view().num_rows() > 0: - await ch_out.send(context, Message(pid, chunk)) + await ch_out.send(context, Message(pid, chunk)) await ch_out.drain(context) diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index 557b6634ebf9..f3d14edf939e 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -35,24 +35,32 @@ _SCHEMA = {"key": DataType(pl.Int32()), "val": DataType(pl.Int32())} _NAMES = list(_SCHEMA) _DTYPES = list(_SCHEMA.values()) +_Boundary = int | tuple[int, ...] + + +def _boundary_value(boundary: _Boundary, index: int) -> int: + return boundary[index] if isinstance(boundary, tuple) else boundary def _make_scheme( context, - boundary: int | tuple[int, ...], + boundary: _Boundary | list[_Boundary], *, key_indices: tuple[int, ...] = (0,), strict: bool = True, stream, ) -> OrderScheme: - boundary_values = ( - boundary if isinstance(boundary, tuple) else (boundary,) * len(key_indices) + boundary_rows: list[_Boundary] = ( + boundary if isinstance(boundary, list) else [boundary] ) boundary_df = DataFrame.from_polars( pl.DataFrame( { - f"k{i}": pl.Series([value], dtype=pl.Int32()) - for i, value in enumerate(boundary_values) + f"k{i}": pl.Series( + [_boundary_value(value, i) for value in boundary_rows], + dtype=pl.Int32(), + ) + for i in range(len(key_indices)) } ), stream, @@ -278,12 +286,46 @@ def test_adjust_orderscheme_sparse_boundary_shift( assert result["val"].to_list() == expected[comm.rank] +@pytest.mark.spmd +def test_adjust_orderscheme_emits_empty_owned_partitions(spmd_engine) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks != 2: + pytest.skip("This test expects exactly two ranks.") + + keys = [0, 1, 2] if comm.rank == 0 else [5, 8] + stream = context.get_stream_from_pool() + input_scheme = _make_scheme(context, 5, stream=stream) + output_scheme = _make_scheme(context, [3, 5, 7], stream=stream) + + with reserve_op_id() as op_id: + output = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(keys), + input_scheme, + output_scheme, + collective_id=op_id, + ) + ) + + expected = { + 0: {0: [0, 1, 2], 1: []}, + 1: {2: [5], 3: [8]}, + }[comm.rank] + assert set(output) == set(expected) + for pid, keys in expected.items(): + assert output[pid]["key"].to_list() == keys + assert output[pid]["val"].to_list() == keys + + @pytest.mark.spmd @pytest.mark.parametrize( "target_boundary,expected", [ (3, {0: [0, 1, 2], 1: [3, 4, 5, 6, 7]}), - (0, {1: list(range(8))}), + (0, {0: [], 1: list(range(8))}), ], ) def test_adjust_orderscheme_single_rank_no_collective( From e8a08340d2add40262a55d4336f61cb37f7dfcf0 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 1 Jun 2026 10:10:37 -0700 Subject: [PATCH 03/26] cleanup --- .../actor_graph/collectives/orderscheme.py | 52 +++++++------------ .../streaming/test_adjust_orderscheme.py | 24 +-------- 2 files changed, 20 insertions(+), 56 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py index ed971f540072..6e7e76adbaf9 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -23,8 +23,6 @@ from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: - from collections.abc import Iterable - from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.streaming.core.channel import Channel @@ -105,15 +103,6 @@ def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Tabl return plc.Table([*table.columns(), pid_col]) -def _pack_table(table: plc.Table, stream: Stream, br: BufferResource) -> PackedData: - """Pack a pylibcudf table as a RapidsMPF PackedData payload.""" - return PackedData.from_cudf_packed_columns( - pack(table, stream, mr=br.device_mr), - stream, - br, - ) - - def _unpack_remote_piece( packed: PackedData, stream: Stream, @@ -145,23 +134,18 @@ def _unpack_remote_piece( ) -def _materialize_packed_pieces( - pieces: Iterable[PackedData], - context: Context, +def _copy_to_owned_chunk( + table: plc.Table, stream: Stream, -) -> TableChunk | None: - """Materialize packed pieces into a uniquely-owned chunk.""" - pieces = list(pieces) - if not pieces: - return None - table = unpack_and_concat(pieces, stream=stream, br=context.br()) - if table.num_rows() == 0: - return None + br: BufferResource, +) -> TableChunk: + """Copy a table view into a uniquely-owned chunk.""" + table = plc.concatenate.concatenate([table], stream=stream, mr=br.device_mr) return TableChunk.from_pylibcudf_table( table, stream, exclusive_view=True, - br=context.br(), + br=br, ) @@ -223,7 +207,7 @@ async def adjust_orderscheme( ) boundary_chunk = output_scheme.get_boundaries(context.br()) boundary_table = boundary_chunk.table_view() - local_pieces: dict[int, list[PackedData]] = defaultdict(list) + local_chunks: dict[int, list[TableChunk]] = defaultdict(list) while (msg := await ch_in.recv(context)) is not None: chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill( @@ -244,13 +228,19 @@ async def adjust_orderscheme( continue owner = _contiguous_owner(pid, comm.nranks, npartitions) if owner == comm.rank: - local_pieces[pid].append(_pack_table(piece, stream, context.br())) + local_chunks[pid].append( + _copy_to_owned_chunk(piece, stream, context.br()) + ) else: assert exchange is not None exchange.insert( owner, - _pack_table( - _append_partition_id(piece, pid, stream), + PackedData.from_cudf_packed_columns( + pack( + _append_partition_id(piece, pid, stream), + stream, + mr=context.br().device_mr, + ), stream, context.br(), ), @@ -262,12 +252,8 @@ async def adjust_orderscheme( output_chunks: dict[int, list[TableChunk]] = defaultdict(list) for source_rank in range(comm.nranks): if source_rank == comm.rank: - for pid, pieces in local_pieces.items(): - chunk = _materialize_packed_pieces( - pieces, context, context.get_stream_from_pool() - ) - if chunk is not None: - output_chunks[pid].append(chunk) + for pid, chunks in local_chunks.items(): + output_chunks[pid].extend(chunks) continue assert exchange is not None stream = context.get_stream_from_pool() diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index f3d14edf939e..1d1a3488cb75 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -9,10 +9,8 @@ import pytest from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.cudf.channel_metadata import ( - ChannelMetadata, OrderKey, OrderScheme, - Partitioning, ) from rapidsmpf.streaming.cudf.table_chunk import TableChunk @@ -26,11 +24,7 @@ from cudf_polars.streaming.actor_graph.collectives.orderscheme import ( adjust_orderscheme, ) -from cudf_polars.streaming.actor_graph.utils import ( - gather_in_task_group, - recv_metadata, - send_metadata, -) +from cudf_polars.streaming.actor_graph.utils import gather_in_task_group _SCHEMA = {"key": DataType(pl.Int32()), "val": DataType(pl.Int32())} _NAMES = list(_SCHEMA) @@ -84,13 +78,6 @@ def _ref_ir() -> Empty: return Empty(_SCHEMA) -def _local_count(comm, scheme: OrderScheme) -> int: - npartitions = scheme.num_boundaries + 1 - return sum( - pid * comm.nranks // npartitions == comm.rank for pid in range(npartitions) - ) - - def _frame(values: list[int]) -> pl.DataFrame: return pl.DataFrame( { @@ -125,14 +112,6 @@ async def _adjust_and_collect( async def _produce() -> None: df = DataFrame.from_polars(input_df, stream) - await send_metadata( - ch_out, - context, - ChannelMetadata( - local_count=_local_count(comm, output_scheme), - partitioning=Partitioning(output_scheme, "inherit"), - ), - ) await ch_in.send( context, Message( @@ -148,7 +127,6 @@ async def _produce() -> None: await ch_in.drain(context) async def _consume() -> None: - await recv_metadata(ch_out, context) while (msg := await ch_out.recv(context)) is not None: output[msg.sequence_number] = _chunk_to_polars( TableChunk.from_message(msg, br=context.br()) From 1f6697e9a68e6cfef00009927a4dc3ab81970d6b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 1 Jun 2026 10:36:52 -0700 Subject: [PATCH 04/26] partial cleanup --- .../actor_graph/collectives/orderscheme.py | 3 +- .../streaming/test_adjust_orderscheme.py | 127 +++++++++++++----- 2 files changed, 97 insertions(+), 33 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py index 6e7e76adbaf9..ea8cd009da90 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -189,7 +189,8 @@ async def adjust_orderscheme( ----- This utility is intentionally narrow and only adjusts data messages. The caller is responsible for receiving input metadata and sending output - metadata. + metadata. Input rows are assumed to be globally ordered by ``input_scheme``; + sortedness is not checked here. """ _validate_schemes(input_scheme, output_scheme) npartitions = output_scheme.num_boundaries + 1 diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index 1d1a3488cb75..354a9714243d 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -74,10 +74,6 @@ def _make_scheme( ) -def _ref_ir() -> Empty: - return Empty(_SCHEMA) - - def _frame(values: list[int]) -> pl.DataFrame: return pl.DataFrame( { @@ -99,31 +95,34 @@ def _chunk_to_polars(chunk: TableChunk) -> pl.DataFrame: async def _adjust_and_collect( context, comm, - input_df: pl.DataFrame, + input_df: pl.DataFrame | list[pl.DataFrame], input_scheme: OrderScheme, output_scheme: OrderScheme, *, collective_id: int | None = None, ) -> dict[int, pl.DataFrame]: + """Run adjustment and collect output chunks by partition ID.""" ch_in = context.create_channel() ch_out = context.create_channel() stream = context.get_stream_from_pool() output: dict[int, pl.DataFrame] = {} async def _produce() -> None: - df = DataFrame.from_polars(input_df, stream) - await ch_in.send( - context, - Message( - comm.rank, - TableChunk.from_pylibcudf_table( - df.table, - stream, - exclusive_view=True, - br=context.br(), + input_dfs = input_df if isinstance(input_df, list) else [input_df] + for sequence_number, df in enumerate(input_dfs): + cudf_df = DataFrame.from_polars(df, stream) + await ch_in.send( + context, + Message( + sequence_number, + TableChunk.from_pylibcudf_table( + cudf_df.table, + stream, + exclusive_view=True, + br=context.br(), + ), ), - ), - ) + ) await ch_in.drain(context) async def _consume() -> None: @@ -141,7 +140,7 @@ async def _consume() -> None: adjust_orderscheme( context, comm, - _ref_ir(), + Empty(_SCHEMA), ir_context, ch_out, ch_in, @@ -155,6 +154,16 @@ async def _consume() -> None: return output +def _assert_partition_output( + output: dict[int, pl.DataFrame], expected: dict[int, list[int]] +) -> None: + """Assert output partition IDs and per-partition row order.""" + assert set(output) == set(expected) + for pid, keys in expected.items(): + assert output[pid]["key"].to_list() == keys + assert output[pid]["val"].to_list() == keys + + async def _adjust_direct( context, comm, @@ -172,7 +181,7 @@ async def _adjust_direct( await adjust_orderscheme( context, comm, - _ref_ir(), + Empty(_SCHEMA), ir_context, ch_out, ch_in, @@ -229,8 +238,8 @@ def test_adjust_orderscheme_requires_collective_id(spmd_engine) -> None: @pytest.mark.parametrize( "target_boundary,expected", [ - (3, {0: [0, 1, 2], 1: [3, 4, 5, 6, 7]}), - (5, {0: [0, 1, 2, 3, 4], 1: [5, 6, 7]}), + (3, {0: {0: [0, 1, 2]}, 1: {1: [3, 4, 5, 6, 7]}}), + (5, {0: {0: [0, 1, 2, 3, 4]}, 1: {1: [5, 6, 7]}}), ], ) def test_adjust_orderscheme_sparse_boundary_shift( @@ -259,9 +268,7 @@ def test_adjust_orderscheme_sparse_boundary_shift( ) ) - result = pl.concat(output.values()) - assert result["key"].to_list() == expected[comm.rank] - assert result["val"].to_list() == expected[comm.rank] + _assert_partition_output(output, expected[comm.rank]) @pytest.mark.spmd @@ -292,10 +299,46 @@ def test_adjust_orderscheme_emits_empty_owned_partitions(spmd_engine) -> None: 0: {0: [0, 1, 2], 1: []}, 1: {2: [5], 3: [8]}, }[comm.rank] - assert set(output) == set(expected) - for pid, keys in expected.items(): - assert output[pid]["key"].to_list() == keys - assert output[pid]["val"].to_list() == keys + _assert_partition_output(output, expected) + + +@pytest.mark.spmd +def test_adjust_orderscheme_all_empty_input(spmd_engine) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + stream = context.get_stream_from_pool() + input_scheme = _make_scheme(context, 5, stream=stream) + output_scheme = _make_scheme(context, [3, 5, 7], stream=stream) + expected = { + pid: [] + for pid in range(output_scheme.num_boundaries + 1) + if pid * comm.nranks // (output_scheme.num_boundaries + 1) == comm.rank + } + + if comm.nranks == 1: + output = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame([]), + input_scheme, + output_scheme, + ) + ) + else: + with reserve_op_id() as op_id: + output = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame([]), + input_scheme, + output_scheme, + collective_id=op_id, + ) + ) + + _assert_partition_output(output, expected) @pytest.mark.spmd @@ -327,7 +370,27 @@ def test_adjust_orderscheme_single_rank_no_collective( ) ) - assert set(output_by_pid) == set(expected) - for pid, keys in expected.items(): - assert output_by_pid[pid]["key"].to_list() == keys - assert output_by_pid[pid]["val"].to_list() == keys + _assert_partition_output(output_by_pid, expected) + + +@pytest.mark.spmd +def test_adjust_orderscheme_multi_chunk_input(spmd_engine) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks != 1: + pytest.skip("This test covers local chunk accumulation.") + + stream = context.get_stream_from_pool() + input_scheme = _make_scheme(context, 4, stream=stream) + output_scheme = _make_scheme(context, 4, stream=stream) + output = asyncio.run( + _adjust_and_collect( + context, + comm, + [_frame([0, 1]), _frame([2, 3, 4, 5]), _frame([6, 7])], + input_scheme, + output_scheme, + ) + ) + + _assert_partition_output(output, {0: [0, 1, 2, 3], 1: [4, 5, 6, 7]}) From 798622a9f29f02720b312e131c485f16e63eb6be Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 1 Jun 2026 11:07:12 -0700 Subject: [PATCH 05/26] avoid all-to-all ctrl msgs --- .../actor_graph/collectives/orderscheme.py | 301 +++++++++++++----- .../streaming/test_adjust_orderscheme.py | 56 +++- 2 files changed, 263 insertions(+), 94 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py index ea8cd009da90..4346d49d13fa 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -19,7 +19,11 @@ from pylibcudf.contiguous_split import pack from cudf_polars.containers import DataFrame, DataType -from cudf_polars.streaming.actor_graph.utils import concat_batch, empty_table_chunk +from cudf_polars.streaming.actor_graph.utils import ( + concat_batch, + empty_table_chunk, + gather_in_task_group, +) from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: @@ -43,13 +47,37 @@ def _contiguous_owner(pid: int, nranks: int, npartitions: int) -> int: return pid * nranks // npartitions +def _partition_range(rank: int, nranks: int, npartitions: int) -> tuple[int, int]: + """Return the half-open partition ID range owned by *rank*.""" + return ( + (rank * npartitions + nranks - 1) // nranks, + ((rank + 1) * npartitions + nranks - 1) // nranks, + ) + + def _local_partitions(rank: int, nranks: int, npartitions: int) -> list[int]: """Return partition IDs owned by *rank* under contiguous assignment.""" - return [ - pid - for pid in range(npartitions) - if _contiguous_owner(pid, nranks, npartitions) == rank - ] + start, stop = _partition_range(rank, nranks, npartitions) + return list(range(start, stop)) + + +def _contiguous_owners( + start: int, + stop: int, + nranks: int, + npartitions: int, +) -> list[int]: + """Return ranks owning any partition in the half-open range [start, stop).""" + if start >= stop: + return [] + first_rank = _contiguous_owner(start, nranks, npartitions) + last_rank = _contiguous_owner(stop - 1, nranks, npartitions) + owners = [] + for rank in range(first_rank, last_rank + 1): + rank_start, rank_stop = _partition_range(rank, nranks, npartitions) + if max(start, rank_start) < min(stop, rank_stop): + owners.append(rank) + return owners def _validate_schemes(input_scheme: OrderScheme, output_scheme: OrderScheme) -> None: @@ -103,6 +131,89 @@ def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Tabl return plc.Table([*table.columns(), pid_col]) +def _boundary_search_positions( + input_boundary_table: plc.Table, + output_boundary_table: plc.Table, + output_scheme: OrderScheme, + stream: Stream, +) -> tuple[list[int], list[int]]: + """Search output boundary positions for projected input boundary rows.""" + if input_boundary_table.num_rows() == 0: + return [], [] + prefix_len = len(output_scheme.keys) + input_prefix_boundaries = plc.Table(input_boundary_table.columns()[:prefix_len]) + orders = [key.order for key in output_scheme.keys] + null_orders = [key.null_order for key in output_scheme.keys] + lower_col = plc.search.lower_bound( + output_boundary_table, + input_prefix_boundaries, + orders, + null_orders, + stream=stream, + ) + upper_col = plc.search.upper_bound( + output_boundary_table, + input_prefix_boundaries, + orders, + null_orders, + stream=stream, + ) + positions = DataFrame.from_table( + plc.Table([lower_col, upper_col]), + ["lower", "upper"], + [_PID_DTYPE, _PID_DTYPE], + stream, + ).to_polars() + return positions["lower"].to_list(), positions["upper"].to_list() + + +def _peer_ranks( + rank: int, + nranks: int, + input_scheme: OrderScheme, + output_scheme: OrderScheme, + lower_positions: list[int], + upper_positions: list[int], +) -> tuple[list[int], list[int]]: + """Return source and destination ranks needed for OrderScheme adjustment.""" + input_npartitions = input_scheme.num_boundaries + 1 + output_npartitions = output_scheme.num_boundaries + 1 + output_prefix_only = len(output_scheme.keys) < len(input_scheme.keys) + include_upper_boundary = output_prefix_only or not input_scheme.strict_boundaries + + def dsts_for_source(source_rank: int) -> list[int]: + input_start, input_stop = _partition_range( + source_rank, nranks, input_npartitions + ) + if input_start == input_stop: + return [] + output_start = 0 if input_start == 0 else upper_positions[input_start - 1] + output_stop = ( + output_npartitions + if input_stop == input_npartitions + else ( + upper_positions[input_stop - 1] + 1 + if include_upper_boundary + else lower_positions[input_stop - 1] + 1 + ) + ) + return [ + dst + for dst in _contiguous_owners( + output_start, output_stop, nranks, output_npartitions + ) + if dst != source_rank + ] + + dsts = dsts_for_source(rank) + srcs = [ + source_rank + for source_rank in range(nranks) + if source_rank != rank and rank in dsts_for_source(source_rank) + ] + return srcs, dsts + + def _unpack_remote_piece( packed: PackedData, stream: Stream, @@ -199,78 +310,112 @@ async def adjust_orderscheme( if comm.nranks > 1 and collective_id is None: raise ValueError("collective_id is required when comm.nranks > 1.") - # TODO: Narrow the peer list (avoid all-to-all control messages) - peers = [rank for rank in range(comm.nranks) if rank != comm.rank] - exchange = ( - SparseAlltoall(context, comm, collective_id, srcs=peers, dsts=peers) - if comm.nranks > 1 - else None - ) - boundary_chunk = output_scheme.get_boundaries(context.br()) - boundary_table = boundary_chunk.table_view() - local_chunks: dict[int, list[TableChunk]] = defaultdict(list) - - while (msg := await ch_in.recv(context)) is not None: - chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill( - context.br(), allow_overbooking=True + try: + input_boundary_chunk = input_scheme.get_boundaries(context.br()) + boundary_chunk = output_scheme.get_boundaries(context.br()) + boundary_table = boundary_chunk.table_view() + srcs: list[int] = [] + dsts: list[int] = [] + if comm.nranks > 1: + with stream_ordered_after( + context.get_stream_from_pool, + upstreams=(input_boundary_chunk.stream, boundary_chunk.stream), + ) as stream: + lower_positions, upper_positions = _boundary_search_positions( + input_boundary_chunk.table_view(), + boundary_table, + output_scheme, + stream, + ) + srcs, dsts = _peer_ranks( + comm.rank, + comm.nranks, + input_scheme, + output_scheme, + lower_positions, + upper_positions, + ) + exchange = ( + SparseAlltoall(context, comm, collective_id, srcs=srcs, dsts=dsts) + if comm.nranks > 1 + else None ) - if chunk.table_view().num_rows() == 0: - continue - with stream_ordered_after( - context.get_stream_from_pool, - upstreams=(chunk.stream, boundary_chunk.stream), - ) as stream: - table = chunk.table_view() - splits = _split_points(table, boundary_table, output_scheme, stream) - for pid, piece in enumerate( - plc.copying.split(table, splits, stream=stream) - ): - if piece.num_rows() == 0: + local_chunks: dict[int, list[TableChunk]] = defaultdict(list) + + try: + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + if chunk.table_view().num_rows() == 0: continue - owner = _contiguous_owner(pid, comm.nranks, npartitions) - if owner == comm.rank: - local_chunks[pid].append( - _copy_to_owned_chunk(piece, stream, context.br()) - ) - else: - assert exchange is not None - exchange.insert( - owner, - PackedData.from_cudf_packed_columns( - pack( - _append_partition_id(piece, pid, stream), - stream, - mr=context.br().device_mr, - ), - stream, - context.br(), - ), - ) - - if exchange is not None: - await exchange.insert_finished(context) - - output_chunks: dict[int, list[TableChunk]] = defaultdict(list) - for source_rank in range(comm.nranks): - if source_rank == comm.rank: - for pid, chunks in local_chunks.items(): - output_chunks[pid].extend(chunks) - continue - assert exchange is not None - stream = context.get_stream_from_pool() - for packed in exchange.extract(source_rank): - remote_piece = _unpack_remote_piece(packed, stream, context.br()) - if remote_piece is None: - continue - pid, chunk = remote_piece - output_chunks[pid].append(chunk) - - for pid in local_pids: - chunks = output_chunks[pid] - chunk = ( - await concat_batch(chunks, context, ref_ir.schema, ir_context) - if chunks - else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) + with stream_ordered_after( + context.get_stream_from_pool, + upstreams=(chunk.stream, boundary_chunk.stream), + ) as stream: + table = chunk.table_view() + splits = _split_points(table, boundary_table, output_scheme, stream) + for pid, piece in enumerate( + plc.copying.split(table, splits, stream=stream) + ): + if piece.num_rows() == 0: + continue + owner = _contiguous_owner(pid, comm.nranks, npartitions) + if owner == comm.rank: + local_chunks[pid].append( + _copy_to_owned_chunk(piece, stream, context.br()) + ) + else: + assert exchange is not None + exchange.insert( + owner, + PackedData.from_cudf_packed_columns( + pack( + _append_partition_id(piece, pid, stream), + stream, + mr=context.br().device_mr, + ), + stream, + context.br(), + ), + ) + finally: + if exchange is not None: + await exchange.insert_finished(context) + + output_chunks: dict[int, list[TableChunk]] = defaultdict(list) + for source_rank in ( + *[src for src in srcs if src < comm.rank], + comm.rank, + *[src for src in srcs if src > comm.rank], + ): + if source_rank == comm.rank: + for pid, chunks in local_chunks.items(): + output_chunks[pid].extend(chunks) + else: + assert exchange is not None + stream = context.get_stream_from_pool() + for packed in exchange.extract(source_rank): + remote_piece = _unpack_remote_piece(packed, stream, context.br()) + if remote_piece is None: + continue + pid, chunk = remote_piece + output_chunks[pid].append(chunk) + + for pid in local_pids: + chunks = output_chunks[pid] + chunk = ( + await concat_batch(chunks, context, ref_ir.schema, ir_context) + if chunks + else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) + ) + await ch_out.send(context, Message(pid, chunk)) + await ch_out.drain(context) + except BaseException: + await gather_in_task_group( + ch_in.shutdown(context), + ch_in.shutdown_metadata(context), + ch_out.shutdown(context), + ch_out.shutdown_metadata(context), ) - await ch_out.send(context, Message(pid, chunk)) - await ch_out.drain(context) + raise diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index 354a9714243d..8652b4308000 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -5,6 +5,7 @@ import asyncio from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING import pytest from rapidsmpf.streaming.core.message import Message @@ -26,10 +27,20 @@ ) from cudf_polars.streaming.actor_graph.utils import gather_in_task_group +if TYPE_CHECKING: + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.context import Context + + from rmm.pylibrmm.stream import Stream + + from cudf_polars.engine.spmd import SPMDEngine + _SCHEMA = {"key": DataType(pl.Int32()), "val": DataType(pl.Int32())} _NAMES = list(_SCHEMA) _DTYPES = list(_SCHEMA.values()) _Boundary = int | tuple[int, ...] +_ExpectedPartitions = dict[int, list[int]] +_ExpectedByRank = dict[int, _ExpectedPartitions] def _boundary_value(boundary: _Boundary, index: int) -> int: @@ -37,12 +48,12 @@ def _boundary_value(boundary: _Boundary, index: int) -> int: def _make_scheme( - context, + context: Context, boundary: _Boundary | list[_Boundary], *, key_indices: tuple[int, ...] = (0,), strict: bool = True, - stream, + stream: Stream, ) -> OrderScheme: boundary_rows: list[_Boundary] = ( boundary if isinstance(boundary, list) else [boundary] @@ -93,8 +104,8 @@ def _chunk_to_polars(chunk: TableChunk) -> pl.DataFrame: async def _adjust_and_collect( - context, - comm, + context: Context, + comm: Communicator, input_df: pl.DataFrame | list[pl.DataFrame], input_scheme: OrderScheme, output_scheme: OrderScheme, @@ -165,8 +176,8 @@ def _assert_partition_output( async def _adjust_direct( - context, - comm, + context: Context, + comm: Communicator, input_scheme: OrderScheme, output_scheme: OrderScheme, *, @@ -193,14 +204,19 @@ async def _adjust_direct( @pytest.mark.spmd @pytest.mark.parametrize( - "input_keys,output_keys,strict,error,match", + "input_keys,output_keys,strict,err,match", [ ((1,), (0,), True, NotImplementedError, "prefix"), ((0,), (0,), False, ValueError, "strict output"), ], ) def test_adjust_orderscheme_rejects_invalid_schemes( - spmd_engine, input_keys, output_keys, strict, error, match + spmd_engine: SPMDEngine, + input_keys: tuple[int, ...], + output_keys: tuple[int, ...], + strict: bool, # noqa: FBT001 + err: type[Exception], + match: str, ) -> None: context = spmd_engine.context stream = context.get_stream_from_pool() @@ -213,14 +229,16 @@ def test_adjust_orderscheme_rejects_invalid_schemes( stream=stream, ) - with pytest.raises(error, match=match): + with pytest.raises(err, match=match): asyncio.run( _adjust_direct(context, spmd_engine.comm, input_scheme, output_scheme) ) @pytest.mark.spmd -def test_adjust_orderscheme_requires_collective_id(spmd_engine) -> None: +def test_adjust_orderscheme_requires_collective_id( + spmd_engine: SPMDEngine, +) -> None: context = spmd_engine.context comm = spmd_engine.comm if comm.nranks == 1: @@ -243,7 +261,9 @@ def test_adjust_orderscheme_requires_collective_id(spmd_engine) -> None: ], ) def test_adjust_orderscheme_sparse_boundary_shift( - spmd_engine, target_boundary, expected + spmd_engine: SPMDEngine, + target_boundary: int, + expected: _ExpectedByRank, ) -> None: context = spmd_engine.context comm = spmd_engine.comm @@ -272,7 +292,9 @@ def test_adjust_orderscheme_sparse_boundary_shift( @pytest.mark.spmd -def test_adjust_orderscheme_emits_empty_owned_partitions(spmd_engine) -> None: +def test_adjust_orderscheme_emits_empty_owned_partitions( + spmd_engine: SPMDEngine, +) -> None: context = spmd_engine.context comm = spmd_engine.comm if comm.nranks != 2: @@ -303,13 +325,13 @@ def test_adjust_orderscheme_emits_empty_owned_partitions(spmd_engine) -> None: @pytest.mark.spmd -def test_adjust_orderscheme_all_empty_input(spmd_engine) -> None: +def test_adjust_orderscheme_all_empty_input(spmd_engine: SPMDEngine) -> None: context = spmd_engine.context comm = spmd_engine.comm stream = context.get_stream_from_pool() input_scheme = _make_scheme(context, 5, stream=stream) output_scheme = _make_scheme(context, [3, 5, 7], stream=stream) - expected = { + expected: _ExpectedPartitions = { pid: [] for pid in range(output_scheme.num_boundaries + 1) if pid * comm.nranks // (output_scheme.num_boundaries + 1) == comm.rank @@ -350,7 +372,9 @@ def test_adjust_orderscheme_all_empty_input(spmd_engine) -> None: ], ) def test_adjust_orderscheme_single_rank_no_collective( - spmd_engine, target_boundary, expected + spmd_engine: SPMDEngine, + target_boundary: int, + expected: _ExpectedPartitions, ) -> None: context = spmd_engine.context comm = spmd_engine.comm @@ -374,7 +398,7 @@ def test_adjust_orderscheme_single_rank_no_collective( @pytest.mark.spmd -def test_adjust_orderscheme_multi_chunk_input(spmd_engine) -> None: +def test_adjust_orderscheme_multi_chunk_input(spmd_engine: SPMDEngine) -> None: context = spmd_engine.context comm = spmd_engine.comm if comm.nranks != 1: From 08ff757b5aef2a381d0ed376f1fdf2ac6eb3de12 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 1 Jun 2026 14:10:31 -0700 Subject: [PATCH 06/26] use distinct payload column --- .../tests/streaming/test_adjust_orderscheme.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index 8652b4308000..d49d47d26feb 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -85,11 +85,15 @@ def _make_scheme( ) +def _payload_value(key: int) -> int: + return key * 10 + 1 + + def _frame(values: list[int]) -> pl.DataFrame: return pl.DataFrame( { "key": pl.Series(values, dtype=pl.Int32()), - "val": pl.Series(values, dtype=pl.Int32()), + "val": pl.Series([_payload_value(v) for v in values], dtype=pl.Int32()), } ) @@ -172,7 +176,7 @@ def _assert_partition_output( assert set(output) == set(expected) for pid, keys in expected.items(): assert output[pid]["key"].to_list() == keys - assert output[pid]["val"].to_list() == keys + assert output[pid]["val"].to_list() == [_payload_value(key) for key in keys] async def _adjust_direct( From 105e8360c446fa93b2c42fe1a456cc8c3095eb54 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 2 Jun 2026 10:03:47 -0700 Subject: [PATCH 07/26] use exclusive_view=True --- python/cudf_polars/tests/streaming/test_adjust_orderscheme.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index d49d47d26feb..145967e25197 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -78,7 +78,7 @@ def _make_scheme( TableChunk.from_pylibcudf_table( boundary_df.table, stream, - exclusive_view=False, + exclusive_view=True, br=context.br(), ), strict_boundaries=strict, From 6168ec621745e17473b09d12b91f0567e88bf982 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 5 Jun 2026 12:18:38 -0700 Subject: [PATCH 08/26] add single-rank fast path --- .../actor_graph/collectives/orderscheme.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py index 4346d49d13fa..b06aef1a5bc3 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -20,6 +20,7 @@ from cudf_polars.containers import DataFrame, DataType from cudf_polars.streaming.actor_graph.utils import ( + ChunkStore, concat_batch, empty_table_chunk, gather_in_task_group, @@ -260,6 +261,75 @@ def _copy_to_owned_chunk( ) +async def _adjust_orderscheme_local( + context: Context, + ref_ir: IR, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + output_scheme: OrderScheme, +) -> None: + npartitions = output_scheme.num_boundaries + 1 + boundary_chunk = output_scheme.get_boundaries(context.br()) + boundary_table = boundary_chunk.table_view() + pending_pid: int | None = None + pending_chunks: ChunkStore | None = None + next_pid = 0 + + async def emit_pending(pid: int) -> None: + nonlocal pending_pid, pending_chunks + if pending_pid == pid and pending_chunks is not None: + chunks = [ + TableChunk.from_message(msg, br=context.br()) for msg in pending_chunks + ] + chunk = await concat_batch(chunks, context, ref_ir.schema, ir_context) + pending_pid = None + pending_chunks = None + else: + chunk = empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) + await ch_out.send(context, Message(pid, chunk)) + + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill( + context.br(), allow_overbooking=True + ) + if chunk.table_view().num_rows() == 0: + continue + with stream_ordered_after( + context.get_stream_from_pool, + upstreams=(chunk.stream, boundary_chunk.stream), + ) as stream: + table = chunk.table_view() + splits = _split_points(table, boundary_table, output_scheme, stream) + for pid, piece in enumerate( + plc.copying.split(table, splits, stream=stream) + ): + if piece.num_rows() == 0: + continue + if pending_pid is not None and pending_pid != pid: + emitted_pid = pending_pid + await emit_pending(emitted_pid) + next_pid = emitted_pid + 1 + while next_pid < pid: + await emit_pending(next_pid) + next_pid += 1 + if pending_pid is None: + pending_pid = pid + pending_chunks = ChunkStore(context) + assert pending_chunks is not None + pending_chunks.insert( + Message( + pid, + _copy_to_owned_chunk(piece, stream, context.br()), + ) + ) + + while next_pid < npartitions: + await emit_pending(next_pid) + next_pid += 1 + await ch_out.drain(context) + + async def adjust_orderscheme( context: Context, comm: Communicator, @@ -311,6 +381,17 @@ async def adjust_orderscheme( raise ValueError("collective_id is required when comm.nranks > 1.") try: + if comm.nranks == 1: + await _adjust_orderscheme_local( + context, + ref_ir, + ir_context, + ch_out, + ch_in, + output_scheme, + ) + return + input_boundary_chunk = input_scheme.get_boundaries(context.br()) boundary_chunk = output_scheme.get_boundaries(context.br()) boundary_table = boundary_chunk.table_view() From 997176d0e4bc4619429a95fee3e2ef0e0e03c0b7 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 10 Jun 2026 09:29:00 -0700 Subject: [PATCH 09/26] formatting --- .../streaming/actor_graph/collectives/orderscheme.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py index b06aef1a5bc3..dde0931b0804 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -7,17 +7,16 @@ from collections import defaultdict from typing import TYPE_CHECKING +import polars as pl + +import pylibcudf as plc +from pylibcudf.contiguous_split import pack from rapidsmpf.integrations.cudf.partition import unpack_and_concat from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.cudf.table_chunk import TableChunk -import polars as pl - -import pylibcudf as plc -from pylibcudf.contiguous_split import pack - from cudf_polars.containers import DataFrame, DataType from cudf_polars.streaming.actor_graph.utils import ( ChunkStore, @@ -33,7 +32,6 @@ from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context from rapidsmpf.streaming.cudf.channel_metadata import OrderScheme - from rmm.pylibrmm.stream import Stream from cudf_polars.dsl.ir import IR, IRExecutionContext From 2fa66eea1e3597eea6cccbd75c99adccf4d1ae6b Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 10 Jun 2026 09:36:48 -0700 Subject: [PATCH 10/26] formatting --- .../tests/streaming/test_adjust_orderscheme.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index 145967e25197..3b52f4d6df1c 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -8,6 +8,10 @@ from typing import TYPE_CHECKING import pytest + +import polars as pl + +import pylibcudf as plc from rapidsmpf.streaming.core.message import Message from rapidsmpf.streaming.cudf.channel_metadata import ( OrderKey, @@ -15,10 +19,6 @@ ) from rapidsmpf.streaming.cudf.table_chunk import TableChunk -import polars as pl - -import pylibcudf as plc - from cudf_polars.containers import DataFrame, DataType from cudf_polars.dsl.ir import Empty, IRExecutionContext from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id @@ -30,7 +30,6 @@ if TYPE_CHECKING: from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.context import Context - from rmm.pylibrmm.stream import Stream from cudf_polars.engine.spmd import SPMDEngine From eaf5812da398002f8ad457e80872bd59758faa72 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 7 Jul 2026 13:03:39 -0700 Subject: [PATCH 11/26] align with main --- .../actor_graph/collectives/orderscheme.py | 86 +++++++++++-------- .../streaming/test_adjust_orderscheme.py | 62 +++++++------ 2 files changed, 85 insertions(+), 63 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py index dde0931b0804..06bb27a188f6 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """OrderScheme adjustment utilities for the RapidsMPF streaming runtime.""" @@ -10,12 +10,12 @@ import polars as pl import pylibcudf as plc +from cudf_streaming.partition_utils import unpack_and_concat +from cudf_streaming.table_chunk import TableChunk from pylibcudf.contiguous_split import pack -from rapidsmpf.integrations.cudf.partition import unpack_and_concat from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall from rapidsmpf.streaming.core.message import Message -from rapidsmpf.streaming.cudf.table_chunk import TableChunk from cudf_polars.containers import DataFrame, DataType from cudf_polars.streaming.actor_graph.utils import ( @@ -27,11 +27,11 @@ from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: + from cudf_streaming.channel_metadata import OrderScheme, Ordering from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context - from rapidsmpf.streaming.cudf.channel_metadata import OrderScheme from rmm.pylibrmm.stream import Stream from cudf_polars.dsl.ir import IR, IRExecutionContext @@ -41,6 +41,12 @@ _PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32) +def _primary_ordering(scheme: OrderScheme) -> Ordering: + """Return the single ordering supported by adjust_orderscheme for now.""" + (ordering,) = scheme.orderings + return ordering + + def _contiguous_owner(pid: int, nranks: int, npartitions: int) -> int: """Return the rank owning *pid* under contiguous partition assignment.""" return pid * nranks // npartitions @@ -81,10 +87,12 @@ def _contiguous_owners( def _validate_schemes(input_scheme: OrderScheme, output_scheme: OrderScheme) -> None: """Validate the first-pass flat OrderScheme adjustment contract.""" - if not output_scheme.strict_boundaries: + input_ordering = _primary_ordering(input_scheme) + output_ordering = _primary_ordering(output_scheme) + if not output_ordering.strict_boundaries: raise ValueError("adjust_orderscheme requires a strict output OrderScheme.") - prefix_len = len(output_scheme.keys) - if input_scheme.keys[:prefix_len] != output_scheme.keys: + prefix_len = len(output_ordering.keys) + if input_ordering.keys[:prefix_len] != output_ordering.keys: raise NotImplementedError( "adjust_orderscheme currently requires the output OrderScheme keys " "to be a prefix of the input OrderScheme keys." @@ -94,18 +102,18 @@ def _validate_schemes(input_scheme: OrderScheme, output_scheme: OrderScheme) -> def _split_points( table: plc.Table, boundary_table: plc.Table, - scheme: OrderScheme, + ordering: Ordering, stream: Stream, ) -> list[int]: """Return row split points that partition *table* by *scheme* boundaries.""" if boundary_table.num_rows() == 0: return [] - key_table = plc.Table([table.columns()[key.column_index] for key in scheme.keys]) + key_table = plc.Table([table.columns()[key.column_index] for key in ordering.keys]) split_col = plc.search.lower_bound( key_table, boundary_table, - [key.order for key in scheme.keys], - [key.null_order for key in scheme.keys], + [key.order for key in ordering.keys], + [key.null_order for key in ordering.keys], stream=stream, ) return ( @@ -133,16 +141,16 @@ def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Tabl def _boundary_search_positions( input_boundary_table: plc.Table, output_boundary_table: plc.Table, - output_scheme: OrderScheme, + output_ordering: Ordering, stream: Stream, ) -> tuple[list[int], list[int]]: """Search output boundary positions for projected input boundary rows.""" if input_boundary_table.num_rows() == 0: return [], [] - prefix_len = len(output_scheme.keys) + prefix_len = len(output_ordering.keys) input_prefix_boundaries = plc.Table(input_boundary_table.columns()[:prefix_len]) - orders = [key.order for key in output_scheme.keys] - null_orders = [key.null_order for key in output_scheme.keys] + orders = [key.order for key in output_ordering.keys] + null_orders = [key.null_order for key in output_ordering.keys] lower_col = plc.search.lower_bound( output_boundary_table, input_prefix_boundaries, @@ -169,16 +177,16 @@ def _boundary_search_positions( def _peer_ranks( rank: int, nranks: int, - input_scheme: OrderScheme, - output_scheme: OrderScheme, + input_ordering: Ordering, + output_ordering: Ordering, lower_positions: list[int], upper_positions: list[int], ) -> tuple[list[int], list[int]]: """Return source and destination ranks needed for OrderScheme adjustment.""" - input_npartitions = input_scheme.num_boundaries + 1 - output_npartitions = output_scheme.num_boundaries + 1 - output_prefix_only = len(output_scheme.keys) < len(input_scheme.keys) - include_upper_boundary = output_prefix_only or not input_scheme.strict_boundaries + input_npartitions = input_ordering.num_boundaries + 1 + output_npartitions = output_ordering.num_boundaries + 1 + output_prefix_only = len(output_ordering.keys) < len(input_ordering.keys) + include_upper_boundary = output_prefix_only or not input_ordering.strict_boundaries def dsts_for_source(source_rank: int) -> list[int]: input_start, input_stop = _partition_range( @@ -265,10 +273,10 @@ async def _adjust_orderscheme_local( ir_context: IRExecutionContext, ch_out: Channel[TableChunk], ch_in: Channel[TableChunk], - output_scheme: OrderScheme, + output_ordering: Ordering, ) -> None: - npartitions = output_scheme.num_boundaries + 1 - boundary_chunk = output_scheme.get_boundaries(context.br()) + npartitions = output_ordering.num_boundaries + 1 + boundary_chunk = output_ordering.get_boundaries(context.br()) boundary_table = boundary_chunk.table_view() pending_pid: int | None = None pending_chunks: ChunkStore | None = None @@ -294,11 +302,11 @@ async def emit_pending(pid: int) -> None: if chunk.table_view().num_rows() == 0: continue with stream_ordered_after( - context.get_stream_from_pool, + context.br().stream_pool.get_stream, upstreams=(chunk.stream, boundary_chunk.stream), ) as stream: table = chunk.table_view() - splits = _split_points(table, boundary_table, output_scheme, stream) + splits = _split_points(table, boundary_table, output_ordering, stream) for pid, piece in enumerate( plc.copying.split(table, splits, stream=stream) ): @@ -372,7 +380,9 @@ async def adjust_orderscheme( sortedness is not checked here. """ _validate_schemes(input_scheme, output_scheme) - npartitions = output_scheme.num_boundaries + 1 + input_ordering = _primary_ordering(input_scheme) + output_ordering = _primary_ordering(output_scheme) + npartitions = output_ordering.num_boundaries + 1 local_pids = _local_partitions(comm.rank, comm.nranks, npartitions) if comm.nranks > 1 and collective_id is None: @@ -386,31 +396,31 @@ async def adjust_orderscheme( ir_context, ch_out, ch_in, - output_scheme, + output_ordering, ) return - input_boundary_chunk = input_scheme.get_boundaries(context.br()) - boundary_chunk = output_scheme.get_boundaries(context.br()) + input_boundary_chunk = input_ordering.get_boundaries(context.br()) + boundary_chunk = output_ordering.get_boundaries(context.br()) boundary_table = boundary_chunk.table_view() srcs: list[int] = [] dsts: list[int] = [] if comm.nranks > 1: with stream_ordered_after( - context.get_stream_from_pool, + context.br().stream_pool.get_stream, upstreams=(input_boundary_chunk.stream, boundary_chunk.stream), ) as stream: lower_positions, upper_positions = _boundary_search_positions( input_boundary_chunk.table_view(), boundary_table, - output_scheme, + output_ordering, stream, ) srcs, dsts = _peer_ranks( comm.rank, comm.nranks, - input_scheme, - output_scheme, + input_ordering, + output_ordering, lower_positions, upper_positions, ) @@ -429,11 +439,13 @@ async def adjust_orderscheme( if chunk.table_view().num_rows() == 0: continue with stream_ordered_after( - context.get_stream_from_pool, + context.br().stream_pool.get_stream, upstreams=(chunk.stream, boundary_chunk.stream), ) as stream: table = chunk.table_view() - splits = _split_points(table, boundary_table, output_scheme, stream) + splits = _split_points( + table, boundary_table, output_ordering, stream + ) for pid, piece in enumerate( plc.copying.split(table, splits, stream=stream) ): @@ -473,7 +485,7 @@ async def adjust_orderscheme( output_chunks[pid].extend(chunks) else: assert exchange is not None - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() for packed in exchange.extract(source_rank): remote_piece = _unpack_remote_piece(packed, stream, context.br()) if remote_piece is None: diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index 3b52f4d6df1c..b8dc0b668dd9 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -12,12 +12,13 @@ import polars as pl import pylibcudf as plc -from rapidsmpf.streaming.core.message import Message -from rapidsmpf.streaming.cudf.channel_metadata import ( +from cudf_streaming.channel_metadata import ( OrderKey, OrderScheme, + Ordering, ) -from rapidsmpf.streaming.cudf.table_chunk import TableChunk +from cudf_streaming.table_chunk import TableChunk +from rapidsmpf.streaming.core.message import Message from cudf_polars.containers import DataFrame, DataType from cudf_polars.dsl.ir import Empty, IRExecutionContext @@ -71,16 +72,24 @@ def _make_scheme( ) return OrderScheme( [ - OrderKey(index, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE) - for index in key_indices - ], - TableChunk.from_pylibcudf_table( - boundary_df.table, - stream, - exclusive_view=True, - br=context.br(), - ), - strict_boundaries=strict, + Ordering( + [ + OrderKey( + index, + plc.types.Order.ASCENDING, + plc.types.NullOrder.BEFORE, + ) + for index in key_indices + ], + TableChunk.from_pylibcudf_table( + boundary_df.table, + stream, + exclusive_view=True, + br=context.br(), + ), + strict_boundaries=strict, + ) + ] ) @@ -118,7 +127,7 @@ async def _adjust_and_collect( """Run adjustment and collect output chunks by partition ID.""" ch_in = context.create_channel() ch_out = context.create_channel() - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() output: dict[int, pl.DataFrame] = {} async def _produce() -> None: @@ -147,7 +156,7 @@ async def _consume() -> None: with ThreadPoolExecutor(max_workers=1) as executor: ir_context = IRExecutionContext( - executor, get_cuda_stream=context.get_stream_from_pool + executor, get_cuda_stream=context.br().stream_pool.get_stream ) await gather_in_task_group( _produce(), @@ -190,7 +199,7 @@ async def _adjust_direct( ch_out = context.create_channel() with ThreadPoolExecutor(max_workers=1) as executor: ir_context = IRExecutionContext( - executor, get_cuda_stream=context.get_stream_from_pool + executor, get_cuda_stream=context.br().stream_pool.get_stream ) await adjust_orderscheme( context, @@ -222,7 +231,7 @@ def test_adjust_orderscheme_rejects_invalid_schemes( match: str, ) -> None: context = spmd_engine.context - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() input_scheme = _make_scheme(context, 4, key_indices=input_keys, stream=stream) output_scheme = _make_scheme( context, @@ -247,7 +256,7 @@ def test_adjust_orderscheme_requires_collective_id( if comm.nranks == 1: pytest.skip("collective_id is only required for multi-rank runs.") - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() input_scheme = _make_scheme(context, 4, stream=stream) output_scheme = _make_scheme(context, 4, stream=stream) @@ -274,7 +283,7 @@ def test_adjust_orderscheme_sparse_boundary_shift( pytest.skip("This test expects exactly two ranks.") keys = list(range(4)) if comm.rank == 0 else list(range(4, 8)) - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() # Input sorted on (key, val) is also sorted on the target key prefix. input_scheme = _make_scheme(context, (4, 4), key_indices=(0, 1), stream=stream) output_scheme = _make_scheme(context, target_boundary, stream=stream) @@ -304,7 +313,7 @@ def test_adjust_orderscheme_emits_empty_owned_partitions( pytest.skip("This test expects exactly two ranks.") keys = [0, 1, 2] if comm.rank == 0 else [5, 8] - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() input_scheme = _make_scheme(context, 5, stream=stream) output_scheme = _make_scheme(context, [3, 5, 7], stream=stream) @@ -331,13 +340,14 @@ def test_adjust_orderscheme_emits_empty_owned_partitions( def test_adjust_orderscheme_all_empty_input(spmd_engine: SPMDEngine) -> None: context = spmd_engine.context comm = spmd_engine.comm - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() input_scheme = _make_scheme(context, 5, stream=stream) output_scheme = _make_scheme(context, [3, 5, 7], stream=stream) + output_npartitions = output_scheme.orderings[0].num_boundaries + 1 expected: _ExpectedPartitions = { pid: [] - for pid in range(output_scheme.num_boundaries + 1) - if pid * comm.nranks // (output_scheme.num_boundaries + 1) == comm.rank + for pid in range(output_npartitions) + if pid * comm.nranks // output_npartitions == comm.rank } if comm.nranks == 1: @@ -384,7 +394,7 @@ def test_adjust_orderscheme_single_rank_no_collective( if comm.nranks != 1: pytest.skip("This test covers the single-rank path.") - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() input_scheme = _make_scheme(context, 4, stream=stream) output_scheme = _make_scheme(context, target_boundary, stream=stream) output_by_pid = asyncio.run( @@ -407,7 +417,7 @@ def test_adjust_orderscheme_multi_chunk_input(spmd_engine: SPMDEngine) -> None: if comm.nranks != 1: pytest.skip("This test covers local chunk accumulation.") - stream = context.get_stream_from_pool() + stream = context.br().stream_pool.get_stream() input_scheme = _make_scheme(context, 4, stream=stream) output_scheme = _make_scheme(context, 4, stream=stream) output = asyncio.run( From 4fd7adf9185f8fdca5ec0c45d6fc0e9834731bd7 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 8 Jul 2026 11:22:11 -0700 Subject: [PATCH 12/26] update algorithm(s) --- .../actor_graph/collectives/orderscheme.py | 487 ++++++++++++------ .../streaming/test_adjust_orderscheme.py | 38 ++ 2 files changed, 365 insertions(+), 160 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py index 06bb27a188f6..1cefc1e4ef92 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py @@ -4,16 +4,17 @@ from __future__ import annotations -from collections import defaultdict from typing import TYPE_CHECKING import polars as pl import pylibcudf as plc -from cudf_streaming.partition_utils import unpack_and_concat +from cudf_streaming.partition_utils import ( + packed_data_from_cudf_packed_columns, + unpack_and_concat, +) from cudf_streaming.table_chunk import TableChunk from pylibcudf.contiguous_split import pack -from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall from rapidsmpf.streaming.core.message import Message @@ -30,6 +31,7 @@ from cudf_streaming.channel_metadata import OrderScheme, Ordering from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.memory.buffer_resource import BufferResource + from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context from rmm.pylibrmm.stream import Stream @@ -60,31 +62,6 @@ def _partition_range(rank: int, nranks: int, npartitions: int) -> tuple[int, int ) -def _local_partitions(rank: int, nranks: int, npartitions: int) -> list[int]: - """Return partition IDs owned by *rank* under contiguous assignment.""" - start, stop = _partition_range(rank, nranks, npartitions) - return list(range(start, stop)) - - -def _contiguous_owners( - start: int, - stop: int, - nranks: int, - npartitions: int, -) -> list[int]: - """Return ranks owning any partition in the half-open range [start, stop).""" - if start >= stop: - return [] - first_rank = _contiguous_owner(start, nranks, npartitions) - last_rank = _contiguous_owner(stop - 1, nranks, npartitions) - owners = [] - for rank in range(first_rank, last_rank + 1): - rank_start, rank_stop = _partition_range(rank, nranks, npartitions) - if max(start, rank_start) < min(stop, rank_stop): - owners.append(rank) - return owners - - def _validate_schemes(input_scheme: OrderScheme, output_scheme: OrderScheme) -> None: """Validate the first-pass flat OrderScheme adjustment contract.""" input_ordering = _primary_ordering(input_scheme) @@ -174,51 +151,38 @@ def _boundary_search_positions( return positions["lower"].to_list(), positions["upper"].to_list() -def _peer_ranks( - rank: int, +def _source_output_range( + source_rank: int, nranks: int, input_ordering: Ordering, output_ordering: Ordering, lower_positions: list[int], upper_positions: list[int], -) -> tuple[list[int], list[int]]: - """Return source and destination ranks needed for OrderScheme adjustment.""" +) -> tuple[int, int]: + """Return the half-open output partition range touched by a source rank.""" input_npartitions = input_ordering.num_boundaries + 1 output_npartitions = output_ordering.num_boundaries + 1 output_prefix_only = len(output_ordering.keys) < len(input_ordering.keys) include_upper_boundary = output_prefix_only or not input_ordering.strict_boundaries - - def dsts_for_source(source_rank: int) -> list[int]: - input_start, input_stop = _partition_range( - source_rank, nranks, input_npartitions - ) - if input_start == input_stop: - return [] - output_start = 0 if input_start == 0 else upper_positions[input_start - 1] - output_stop = ( - output_npartitions - if input_stop == input_npartitions - else ( - upper_positions[input_stop - 1] + 1 - if include_upper_boundary - else lower_positions[input_stop - 1] + 1 - ) + input_start, input_stop = _partition_range(source_rank, nranks, input_npartitions) + if input_start == input_stop: + return 0, 0 + output_start = 0 if input_start == 0 else upper_positions[input_start - 1] + output_stop = ( + output_npartitions + if input_stop == input_npartitions + else ( + upper_positions[input_stop - 1] + 1 + if include_upper_boundary + else lower_positions[input_stop - 1] + 1 ) - return [ - dst - for dst in _contiguous_owners( - output_start, output_stop, nranks, output_npartitions - ) - if dst != source_rank - ] + ) + return output_start, output_stop - dsts = dsts_for_source(rank) - srcs = [ - source_rank - for source_rank in range(nranks) - if source_rank != rank and rank in dsts_for_source(source_rank) - ] - return srcs, dsts + +def _ranges_overlap(left: tuple[int, int], right: tuple[int, int]) -> bool: + """Return whether two half-open integer ranges overlap.""" + return max(left[0], right[0]) < min(left[1], right[1]) def _unpack_remote_piece( @@ -267,6 +231,69 @@ def _copy_to_owned_chunk( ) +class _OutputPieceReader: + """Read input only far enough to materialize requested output windows.""" + + def __init__( + self, + context: Context, + ch_in: Channel[TableChunk], + boundary_chunk: TableChunk, + output_ordering: Ordering, + ) -> None: + self.context = context + self.ch_in = ch_in + self.boundary_chunk = boundary_chunk + self.boundary_table = boundary_chunk.table_view() + self.output_ordering = output_ordering + self.pending: dict[int, ChunkStore] = {} + self.input_done = False + + def _store(self, pid: int, chunk: TableChunk) -> None: + if pid not in self.pending: + self.pending[pid] = ChunkStore(self.context) + self.pending[pid].insert(Message(pid, chunk)) + + def _has_reached(self, stop: int) -> bool: + return any(pid >= stop for pid in self.pending) + + async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: + """Return all locally-read pieces for output pids in ``[start, stop)``.""" + while not self.input_done and not self._has_reached(stop): + msg = await self.ch_in.recv(self.context) + if msg is None: + self.input_done = True + break + chunk = TableChunk.from_message( + msg, br=self.context.br() + ).make_available_and_spill(self.context.br(), allow_overbooking=True) + if chunk.table_view().num_rows() == 0: + continue + with stream_ordered_after( + self.context.br().stream_pool.get_stream, + upstreams=(chunk.stream, self.boundary_chunk.stream), + ) as stream: + table = chunk.table_view() + splits = _split_points( + table, self.boundary_table, self.output_ordering, stream + ) + for pid, piece in enumerate( + plc.copying.split(table, splits, stream=stream) + ): + if piece.num_rows() == 0: + continue + self._store( + pid, _copy_to_owned_chunk(piece, stream, self.context.br()) + ) + + out = { + pid: self.pending.pop(pid) + for pid in list(self.pending) + if start <= pid < stop + } + return dict(sorted(out.items())) + + async def _adjust_orderscheme_local( context: Context, ref_ir: IR, @@ -336,6 +363,223 @@ async def emit_pending(pid: int) -> None: await ch_out.drain(context) +def _store_chunk( + context: Context, + stores: dict[int, ChunkStore], + pid: int, + chunk: TableChunk, +) -> None: + if pid not in stores: + stores[pid] = ChunkStore(context) + stores[pid].insert(Message(pid, chunk)) + + +async def _adjust_orderscheme_rank_hybrid( + context: Context, + comm: Communicator, + ref_ir: IR, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + input_ordering: Ordering, + output_ordering: Ordering, + collective_id: int, + lower_positions: list[int], + upper_positions: list[int], +) -> None: + npartitions = output_ordering.num_boundaries + 1 + boundary_chunk = output_ordering.get_boundaries(context.br()) + boundary_table = boundary_chunk.table_view() + local_window = _partition_range(comm.rank, comm.nranks, npartitions) + source_ranges = [ + _source_output_range( + source_rank, + comm.nranks, + input_ordering, + output_ordering, + lower_positions, + upper_positions, + ) + for source_rank in range(comm.nranks) + ] + local_source_range = source_ranges[comm.rank] + srcs = [ + source_rank + for source_rank, source_range in enumerate(source_ranges) + if source_rank != comm.rank and _ranges_overlap(source_range, local_window) + ] + dsts = [ + output_rank + for output_rank in range(comm.nranks) + if output_rank != comm.rank + and _ranges_overlap( + local_source_range, _partition_range(output_rank, comm.nranks, npartitions) + ) + ] + exchange = SparseAlltoall( + context, + comm, + collective_id, + srcs=srcs, + dsts=dsts, + ) + + async def send_piece(pid: int, chunk: TableChunk) -> None: + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(chunk.stream,), + ) as stream: + exchange.insert( + _contiguous_owner(pid, comm.nranks, npartitions), + packed_data_from_cudf_packed_columns( + pack( + _append_partition_id(chunk.table_view(), pid, stream), + stream, + mr=context.br().device_mr, + ), + stream, + context.br(), + ), + ) + + async def emit(pid: int, store: ChunkStore | None) -> None: + chunks = ( + [TableChunk.from_message(msg, br=context.br()) for msg in store] + if store is not None + else [] + ) + chunk = ( + await concat_batch(chunks, context, ref_ir.schema, ir_context) + if chunks + else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) + ) + await ch_out.send(context, Message(pid, chunk)) + + # Ranks with no incoming dependency can stream local output immediately + # while sending remote-owned pieces as they are encountered. + if not srcs: + pending_pid: int | None = None + pending_chunks: ChunkStore | None = None + next_pid = local_window[0] + + async def emit_pending(pid: int) -> None: + nonlocal pending_pid, pending_chunks + await emit(pid, pending_chunks if pending_pid == pid else None) + if pending_pid == pid: + pending_pid = None + pending_chunks = None + + try: + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + if chunk.table_view().num_rows() == 0: + continue + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(chunk.stream, boundary_chunk.stream), + ) as stream: + table = chunk.table_view() + splits = _split_points( + table, boundary_table, output_ordering, stream + ) + for pid, piece in enumerate( + plc.copying.split(table, splits, stream=stream) + ): + if piece.num_rows() == 0: + continue + owner = _contiguous_owner(pid, comm.nranks, npartitions) + owned = owner == comm.rank + piece_chunk = _copy_to_owned_chunk(piece, stream, context.br()) + if not owned: + await send_piece(pid, piece_chunk) + continue + if pending_pid is not None and pending_pid != pid: + emitted_pid = pending_pid + await emit_pending(emitted_pid) + next_pid = emitted_pid + 1 + while next_pid < pid: + await emit_pending(next_pid) + next_pid += 1 + if pending_pid is None: + pending_pid = pid + pending_chunks = ChunkStore(context) + assert pending_chunks is not None + pending_chunks.insert(Message(pid, piece_chunk)) + finally: + await exchange.insert_finished(context) + + while next_pid < local_window[1]: + await emit_pending(next_pid) + next_pid += 1 + await ch_out.drain(context) + return + + reader = _OutputPieceReader(context, ch_in, boundary_chunk, output_ordering) + local_pieces: dict[int, ChunkStore] = {} + # If a higher rank depends on this rank's data, read far enough to make + # that data available before waiting for lower-rank input. + if dsts: + pieces = await reader.collect_window(*local_source_range) + for pid, store in pieces.items(): + owner = _contiguous_owner(pid, comm.nranks, npartitions) + if owner == comm.rank: + local_pieces[pid] = store + continue + for msg in store: + await send_piece(pid, TableChunk.from_message(msg, br=context.br())) + await exchange.insert_finished(context) + + pieces_by_source: dict[int, dict[int, ChunkStore]] = {} + if local_pieces: + pieces_by_source[comm.rank] = local_pieces + for source_rank in srcs: + remote_pieces: dict[int, ChunkStore] = {} + stream = context.br().stream_pool.get_stream() + for packed in exchange.extract(source_rank): + remote_piece = _unpack_remote_piece(packed, stream, context.br()) + if remote_piece is None: + continue + pid, chunk = remote_piece + _store_chunk(context, remote_pieces, pid, chunk) + pieces_by_source[source_rank] = remote_pieces + + for pid, store in (await reader.collect_window(*local_window)).items(): + if _contiguous_owner(pid, comm.nranks, npartitions) == comm.rank: + if pid not in local_pieces: + local_pieces[pid] = ChunkStore(context) + for msg in store: + local_pieces[pid].insert(msg) + if local_pieces: + pieces_by_source[comm.rank] = local_pieces + + contributing_sources = [ + source_rank + for source_rank, source_range in enumerate(source_ranges) + if _ranges_overlap(source_range, local_window) + ] + for pid in range(*local_window): + chunks: list[TableChunk] = [] + for source_rank in contributing_sources: + stores = pieces_by_source.get(source_rank) + if stores is None: + continue + pid_store = stores.get(pid) + if pid_store is None: + continue + chunks.extend( + TableChunk.from_message(msg, br=context.br()) for msg in pid_store + ) + chunk = ( + await concat_batch(chunks, context, ref_ir.schema, ir_context) + if chunks + else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) + ) + await ch_out.send(context, Message(pid, chunk)) + await ch_out.drain(context) + + async def adjust_orderscheme( context: Context, comm: Communicator, @@ -382,8 +626,6 @@ async def adjust_orderscheme( _validate_schemes(input_scheme, output_scheme) input_ordering = _primary_ordering(input_scheme) output_ordering = _primary_ordering(output_scheme) - npartitions = output_ordering.num_boundaries + 1 - local_pids = _local_partitions(comm.rank, comm.nranks, npartitions) if comm.nranks > 1 and collective_id is None: raise ValueError("collective_id is required when comm.nranks > 1.") @@ -403,105 +645,30 @@ async def adjust_orderscheme( input_boundary_chunk = input_ordering.get_boundaries(context.br()) boundary_chunk = output_ordering.get_boundaries(context.br()) boundary_table = boundary_chunk.table_view() - srcs: list[int] = [] - dsts: list[int] = [] - if comm.nranks > 1: - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(input_boundary_chunk.stream, boundary_chunk.stream), - ) as stream: - lower_positions, upper_positions = _boundary_search_positions( - input_boundary_chunk.table_view(), - boundary_table, - output_ordering, - stream, - ) - srcs, dsts = _peer_ranks( - comm.rank, - comm.nranks, - input_ordering, + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(input_boundary_chunk.stream, boundary_chunk.stream), + ) as stream: + lower_positions, upper_positions = _boundary_search_positions( + input_boundary_chunk.table_view(), + boundary_table, output_ordering, - lower_positions, - upper_positions, + stream, ) - exchange = ( - SparseAlltoall(context, comm, collective_id, srcs=srcs, dsts=dsts) - if comm.nranks > 1 - else None + assert collective_id is not None + await _adjust_orderscheme_rank_hybrid( + context, + comm, + ref_ir, + ir_context, + ch_out, + ch_in, + input_ordering, + output_ordering, + collective_id, + lower_positions, + upper_positions, ) - local_chunks: dict[int, list[TableChunk]] = defaultdict(list) - - try: - while (msg := await ch_in.recv(context)) is not None: - chunk = TableChunk.from_message( - msg, br=context.br() - ).make_available_and_spill(context.br(), allow_overbooking=True) - if chunk.table_view().num_rows() == 0: - continue - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(chunk.stream, boundary_chunk.stream), - ) as stream: - table = chunk.table_view() - splits = _split_points( - table, boundary_table, output_ordering, stream - ) - for pid, piece in enumerate( - plc.copying.split(table, splits, stream=stream) - ): - if piece.num_rows() == 0: - continue - owner = _contiguous_owner(pid, comm.nranks, npartitions) - if owner == comm.rank: - local_chunks[pid].append( - _copy_to_owned_chunk(piece, stream, context.br()) - ) - else: - assert exchange is not None - exchange.insert( - owner, - PackedData.from_cudf_packed_columns( - pack( - _append_partition_id(piece, pid, stream), - stream, - mr=context.br().device_mr, - ), - stream, - context.br(), - ), - ) - finally: - if exchange is not None: - await exchange.insert_finished(context) - - output_chunks: dict[int, list[TableChunk]] = defaultdict(list) - for source_rank in ( - *[src for src in srcs if src < comm.rank], - comm.rank, - *[src for src in srcs if src > comm.rank], - ): - if source_rank == comm.rank: - for pid, chunks in local_chunks.items(): - output_chunks[pid].extend(chunks) - else: - assert exchange is not None - stream = context.br().stream_pool.get_stream() - for packed in exchange.extract(source_rank): - remote_piece = _unpack_remote_piece(packed, stream, context.br()) - if remote_piece is None: - continue - pid, chunk = remote_piece - output_chunks[pid].append(chunk) - - for pid in local_pids: - chunks = output_chunks[pid] - chunk = ( - await concat_batch(chunks, context, ref_ir.schema, ir_context) - if chunks - else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) - ) - await ch_out.send(context, Message(pid, chunk)) - await ch_out.drain(context) except BaseException: await gather_in_task_group( ch_in.shutdown(context), diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py index b8dc0b668dd9..6a5f73feeccf 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py @@ -336,6 +336,44 @@ def test_adjust_orderscheme_emits_empty_owned_partitions( _assert_partition_output(output, expected) +@pytest.mark.spmd +def test_adjust_orderscheme_middle_rank_buffers_only_as_needed( + spmd_engine: SPMDEngine, +) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks != 3: + pytest.skip("This test expects exactly three ranks.") + + keys = { + 0: [0, 5, 9], + 1: [10, 15, 19], + 2: [20, 25], + }[comm.rank] + stream = context.br().stream_pool.get_stream() + input_scheme = _make_scheme(context, [10, 20], stream=stream) + output_scheme = _make_scheme(context, [5, 15], stream=stream) + + with reserve_op_id() as op_id: + output = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(keys), + input_scheme, + output_scheme, + collective_id=op_id, + ) + ) + + expected = { + 0: {0: [0]}, + 1: {1: [5, 9, 10]}, + 2: {2: [15, 19, 20, 25]}, + }[comm.rank] + _assert_partition_output(output, expected) + + @pytest.mark.spmd def test_adjust_orderscheme_all_empty_input(spmd_engine: SPMDEngine) -> None: context = spmd_engine.context From 1ea496853580384f876d4f1dd8baab9f3ea023b1 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 8 Jul 2026 12:41:30 -0700 Subject: [PATCH 13/26] use ordering instaed of orderscheme as the focus --- .../{orderscheme.py => ordering.py} | 54 +++---- ...orderscheme.py => test_adjust_ordering.py} | 143 +++++++++--------- 2 files changed, 91 insertions(+), 106 deletions(-) rename python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/{orderscheme.py => ordering.py} (93%) rename python/cudf_polars/tests/streaming/{test_adjust_orderscheme.py => test_adjust_ordering.py} (76%) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py similarity index 93% rename from python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 1cefc1e4ef92..23a643350891 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/orderscheme.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""OrderScheme adjustment utilities for the RapidsMPF streaming runtime.""" +"""Ordering adjustment utilities for the RapidsMPF streaming runtime.""" from __future__ import annotations @@ -28,7 +28,7 @@ from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: - from cudf_streaming.channel_metadata import OrderScheme, Ordering + from cudf_streaming.channel_metadata import Ordering from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.packed_data import PackedData @@ -43,12 +43,6 @@ _PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32) -def _primary_ordering(scheme: OrderScheme) -> Ordering: - """Return the single ordering supported by adjust_orderscheme for now.""" - (ordering,) = scheme.orderings - return ordering - - def _contiguous_owner(pid: int, nranks: int, npartitions: int) -> int: """Return the rank owning *pid* under contiguous partition assignment.""" return pid * nranks // npartitions @@ -62,17 +56,15 @@ def _partition_range(rank: int, nranks: int, npartitions: int) -> tuple[int, int ) -def _validate_schemes(input_scheme: OrderScheme, output_scheme: OrderScheme) -> None: - """Validate the first-pass flat OrderScheme adjustment contract.""" - input_ordering = _primary_ordering(input_scheme) - output_ordering = _primary_ordering(output_scheme) +def _validate_orderings(input_ordering: Ordering, output_ordering: Ordering) -> None: + """Validate the first-pass flat Ordering adjustment contract.""" if not output_ordering.strict_boundaries: - raise ValueError("adjust_orderscheme requires a strict output OrderScheme.") + raise ValueError("adjust_ordering requires a strict output Ordering.") prefix_len = len(output_ordering.keys) if input_ordering.keys[:prefix_len] != output_ordering.keys: raise NotImplementedError( - "adjust_orderscheme currently requires the output OrderScheme keys " - "to be a prefix of the input OrderScheme keys." + "adjust_ordering currently requires the output Ordering keys " + "to be a prefix of the input Ordering keys." ) @@ -82,7 +74,7 @@ def _split_points( ordering: Ordering, stream: Stream, ) -> list[int]: - """Return row split points that partition *table* by *scheme* boundaries.""" + """Return row split points that partition *table* by *ordering* boundaries.""" if boundary_table.num_rows() == 0: return [] key_table = plc.Table([table.columns()[key.column_index] for key in ordering.keys]) @@ -294,7 +286,7 @@ async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: return dict(sorted(out.items())) -async def _adjust_orderscheme_local( +async def _adjust_ordering_local( context: Context, ref_ir: IR, ir_context: IRExecutionContext, @@ -374,7 +366,7 @@ def _store_chunk( stores[pid].insert(Message(pid, chunk)) -async def _adjust_orderscheme_rank_hybrid( +async def _adjust_ordering_rank_hybrid( context: Context, comm: Communicator, ref_ir: IR, @@ -580,20 +572,20 @@ async def emit_pending(pid: int) -> None: await ch_out.drain(context) -async def adjust_orderscheme( +async def adjust_ordering( context: Context, comm: Communicator, ref_ir: IR, ir_context: IRExecutionContext, ch_out: Channel[TableChunk], ch_in: Channel[TableChunk], - input_scheme: OrderScheme, - output_scheme: OrderScheme, + input_ordering: Ordering, + output_ordering: Ordering, *, collective_id: int | None = None, ) -> None: """ - Adjust flat OrderScheme boundaries using contiguous partition ownership. + Adjust flat Ordering boundaries using contiguous partition ownership. Parameters ---------- @@ -609,10 +601,10 @@ async def adjust_orderscheme( The output channel. ch_in The input channel. - input_scheme - The input OrderScheme. - output_scheme - The output OrderScheme. + input_ordering + The input Ordering. + output_ordering + The output Ordering. collective_id The collective ID to use for SparseAlltoall. @@ -620,19 +612,17 @@ async def adjust_orderscheme( ----- This utility is intentionally narrow and only adjusts data messages. The caller is responsible for receiving input metadata and sending output - metadata. Input rows are assumed to be globally ordered by ``input_scheme``; + metadata. Input rows are assumed to be globally ordered by ``input_ordering``; sortedness is not checked here. """ - _validate_schemes(input_scheme, output_scheme) - input_ordering = _primary_ordering(input_scheme) - output_ordering = _primary_ordering(output_scheme) + _validate_orderings(input_ordering, output_ordering) if comm.nranks > 1 and collective_id is None: raise ValueError("collective_id is required when comm.nranks > 1.") try: if comm.nranks == 1: - await _adjust_orderscheme_local( + await _adjust_ordering_local( context, ref_ir, ir_context, @@ -656,7 +646,7 @@ async def adjust_orderscheme( stream, ) assert collective_id is not None - await _adjust_orderscheme_rank_hybrid( + await _adjust_ordering_rank_hybrid( context, comm, ref_ir, diff --git a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py b/python/cudf_polars/tests/streaming/test_adjust_ordering.py similarity index 76% rename from python/cudf_polars/tests/streaming/test_adjust_orderscheme.py rename to python/cudf_polars/tests/streaming/test_adjust_ordering.py index 6a5f73feeccf..d7746816a215 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_orderscheme.py +++ b/python/cudf_polars/tests/streaming/test_adjust_ordering.py @@ -14,7 +14,6 @@ import pylibcudf as plc from cudf_streaming.channel_metadata import ( OrderKey, - OrderScheme, Ordering, ) from cudf_streaming.table_chunk import TableChunk @@ -23,8 +22,8 @@ from cudf_polars.containers import DataFrame, DataType from cudf_polars.dsl.ir import Empty, IRExecutionContext from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id -from cudf_polars.streaming.actor_graph.collectives.orderscheme import ( - adjust_orderscheme, +from cudf_polars.streaming.actor_graph.collectives.ordering import ( + adjust_ordering, ) from cudf_polars.streaming.actor_graph.utils import gather_in_task_group @@ -47,14 +46,14 @@ def _boundary_value(boundary: _Boundary, index: int) -> int: return boundary[index] if isinstance(boundary, tuple) else boundary -def _make_scheme( +def _make_ordering( context: Context, boundary: _Boundary | list[_Boundary], *, key_indices: tuple[int, ...] = (0,), strict: bool = True, stream: Stream, -) -> OrderScheme: +) -> Ordering: boundary_rows: list[_Boundary] = ( boundary if isinstance(boundary, list) else [boundary] ) @@ -70,26 +69,22 @@ def _make_scheme( ), stream, ) - return OrderScheme( + return Ordering( [ - Ordering( - [ - OrderKey( - index, - plc.types.Order.ASCENDING, - plc.types.NullOrder.BEFORE, - ) - for index in key_indices - ], - TableChunk.from_pylibcudf_table( - boundary_df.table, - stream, - exclusive_view=True, - br=context.br(), - ), - strict_boundaries=strict, + OrderKey( + index, + plc.types.Order.ASCENDING, + plc.types.NullOrder.BEFORE, ) - ] + for index in key_indices + ], + TableChunk.from_pylibcudf_table( + boundary_df.table, + stream, + exclusive_view=True, + br=context.br(), + ), + strict_boundaries=strict, ) @@ -119,8 +114,8 @@ async def _adjust_and_collect( context: Context, comm: Communicator, input_df: pl.DataFrame | list[pl.DataFrame], - input_scheme: OrderScheme, - output_scheme: OrderScheme, + input_ordering: Ordering, + output_ordering: Ordering, *, collective_id: int | None = None, ) -> dict[int, pl.DataFrame]: @@ -160,15 +155,15 @@ async def _consume() -> None: ) await gather_in_task_group( _produce(), - adjust_orderscheme( + adjust_ordering( context, comm, Empty(_SCHEMA), ir_context, ch_out, ch_in, - input_scheme, - output_scheme, + input_ordering, + output_ordering, collective_id=collective_id, ), _consume(), @@ -190,8 +185,8 @@ def _assert_partition_output( async def _adjust_direct( context: Context, comm: Communicator, - input_scheme: OrderScheme, - output_scheme: OrderScheme, + input_ordering: Ordering, + output_ordering: Ordering, *, collective_id: int | None = None, ) -> None: @@ -201,15 +196,15 @@ async def _adjust_direct( ir_context = IRExecutionContext( executor, get_cuda_stream=context.br().stream_pool.get_stream ) - await adjust_orderscheme( + await adjust_ordering( context, comm, Empty(_SCHEMA), ir_context, ch_out, ch_in, - input_scheme, - output_scheme, + input_ordering, + output_ordering, collective_id=collective_id, ) @@ -222,7 +217,7 @@ async def _adjust_direct( ((0,), (0,), False, ValueError, "strict output"), ], ) -def test_adjust_orderscheme_rejects_invalid_schemes( +def test_adjust_ordering_rejects_invalid_orderings( spmd_engine: SPMDEngine, input_keys: tuple[int, ...], output_keys: tuple[int, ...], @@ -232,8 +227,8 @@ def test_adjust_orderscheme_rejects_invalid_schemes( ) -> None: context = spmd_engine.context stream = context.br().stream_pool.get_stream() - input_scheme = _make_scheme(context, 4, key_indices=input_keys, stream=stream) - output_scheme = _make_scheme( + input_ordering = _make_ordering(context, 4, key_indices=input_keys, stream=stream) + output_ordering = _make_ordering( context, 4, key_indices=output_keys, @@ -243,12 +238,12 @@ def test_adjust_orderscheme_rejects_invalid_schemes( with pytest.raises(err, match=match): asyncio.run( - _adjust_direct(context, spmd_engine.comm, input_scheme, output_scheme) + _adjust_direct(context, spmd_engine.comm, input_ordering, output_ordering) ) @pytest.mark.spmd -def test_adjust_orderscheme_requires_collective_id( +def test_adjust_ordering_requires_collective_id( spmd_engine: SPMDEngine, ) -> None: context = spmd_engine.context @@ -257,11 +252,11 @@ def test_adjust_orderscheme_requires_collective_id( pytest.skip("collective_id is only required for multi-rank runs.") stream = context.br().stream_pool.get_stream() - input_scheme = _make_scheme(context, 4, stream=stream) - output_scheme = _make_scheme(context, 4, stream=stream) + input_ordering = _make_ordering(context, 4, stream=stream) + output_ordering = _make_ordering(context, 4, stream=stream) with pytest.raises(ValueError, match="collective_id"): - asyncio.run(_adjust_direct(context, comm, input_scheme, output_scheme)) + asyncio.run(_adjust_direct(context, comm, input_ordering, output_ordering)) @pytest.mark.spmd @@ -272,7 +267,7 @@ def test_adjust_orderscheme_requires_collective_id( (5, {0: {0: [0, 1, 2, 3, 4]}, 1: {1: [5, 6, 7]}}), ], ) -def test_adjust_orderscheme_sparse_boundary_shift( +def test_adjust_ordering_sparse_boundary_shift( spmd_engine: SPMDEngine, target_boundary: int, expected: _ExpectedByRank, @@ -285,8 +280,8 @@ def test_adjust_orderscheme_sparse_boundary_shift( keys = list(range(4)) if comm.rank == 0 else list(range(4, 8)) stream = context.br().stream_pool.get_stream() # Input sorted on (key, val) is also sorted on the target key prefix. - input_scheme = _make_scheme(context, (4, 4), key_indices=(0, 1), stream=stream) - output_scheme = _make_scheme(context, target_boundary, stream=stream) + input_ordering = _make_ordering(context, (4, 4), key_indices=(0, 1), stream=stream) + output_ordering = _make_ordering(context, target_boundary, stream=stream) with reserve_op_id() as op_id: output = asyncio.run( @@ -294,8 +289,8 @@ def test_adjust_orderscheme_sparse_boundary_shift( context, comm, _frame(keys), - input_scheme, - output_scheme, + input_ordering, + output_ordering, collective_id=op_id, ) ) @@ -304,7 +299,7 @@ def test_adjust_orderscheme_sparse_boundary_shift( @pytest.mark.spmd -def test_adjust_orderscheme_emits_empty_owned_partitions( +def test_adjust_ordering_emits_empty_owned_partitions( spmd_engine: SPMDEngine, ) -> None: context = spmd_engine.context @@ -314,8 +309,8 @@ def test_adjust_orderscheme_emits_empty_owned_partitions( keys = [0, 1, 2] if comm.rank == 0 else [5, 8] stream = context.br().stream_pool.get_stream() - input_scheme = _make_scheme(context, 5, stream=stream) - output_scheme = _make_scheme(context, [3, 5, 7], stream=stream) + input_ordering = _make_ordering(context, 5, stream=stream) + output_ordering = _make_ordering(context, [3, 5, 7], stream=stream) with reserve_op_id() as op_id: output = asyncio.run( @@ -323,8 +318,8 @@ def test_adjust_orderscheme_emits_empty_owned_partitions( context, comm, _frame(keys), - input_scheme, - output_scheme, + input_ordering, + output_ordering, collective_id=op_id, ) ) @@ -337,7 +332,7 @@ def test_adjust_orderscheme_emits_empty_owned_partitions( @pytest.mark.spmd -def test_adjust_orderscheme_middle_rank_buffers_only_as_needed( +def test_adjust_ordering_middle_rank_buffers_only_as_needed( spmd_engine: SPMDEngine, ) -> None: context = spmd_engine.context @@ -351,8 +346,8 @@ def test_adjust_orderscheme_middle_rank_buffers_only_as_needed( 2: [20, 25], }[comm.rank] stream = context.br().stream_pool.get_stream() - input_scheme = _make_scheme(context, [10, 20], stream=stream) - output_scheme = _make_scheme(context, [5, 15], stream=stream) + input_ordering = _make_ordering(context, [10, 20], stream=stream) + output_ordering = _make_ordering(context, [5, 15], stream=stream) with reserve_op_id() as op_id: output = asyncio.run( @@ -360,8 +355,8 @@ def test_adjust_orderscheme_middle_rank_buffers_only_as_needed( context, comm, _frame(keys), - input_scheme, - output_scheme, + input_ordering, + output_ordering, collective_id=op_id, ) ) @@ -375,13 +370,13 @@ def test_adjust_orderscheme_middle_rank_buffers_only_as_needed( @pytest.mark.spmd -def test_adjust_orderscheme_all_empty_input(spmd_engine: SPMDEngine) -> None: +def test_adjust_ordering_all_empty_input(spmd_engine: SPMDEngine) -> None: context = spmd_engine.context comm = spmd_engine.comm stream = context.br().stream_pool.get_stream() - input_scheme = _make_scheme(context, 5, stream=stream) - output_scheme = _make_scheme(context, [3, 5, 7], stream=stream) - output_npartitions = output_scheme.orderings[0].num_boundaries + 1 + input_ordering = _make_ordering(context, 5, stream=stream) + output_ordering = _make_ordering(context, [3, 5, 7], stream=stream) + output_npartitions = output_ordering.num_boundaries + 1 expected: _ExpectedPartitions = { pid: [] for pid in range(output_npartitions) @@ -394,8 +389,8 @@ def test_adjust_orderscheme_all_empty_input(spmd_engine: SPMDEngine) -> None: context, comm, _frame([]), - input_scheme, - output_scheme, + input_ordering, + output_ordering, ) ) else: @@ -405,8 +400,8 @@ def test_adjust_orderscheme_all_empty_input(spmd_engine: SPMDEngine) -> None: context, comm, _frame([]), - input_scheme, - output_scheme, + input_ordering, + output_ordering, collective_id=op_id, ) ) @@ -422,7 +417,7 @@ def test_adjust_orderscheme_all_empty_input(spmd_engine: SPMDEngine) -> None: (0, {0: [], 1: list(range(8))}), ], ) -def test_adjust_orderscheme_single_rank_no_collective( +def test_adjust_ordering_single_rank_no_collective( spmd_engine: SPMDEngine, target_boundary: int, expected: _ExpectedPartitions, @@ -433,15 +428,15 @@ def test_adjust_orderscheme_single_rank_no_collective( pytest.skip("This test covers the single-rank path.") stream = context.br().stream_pool.get_stream() - input_scheme = _make_scheme(context, 4, stream=stream) - output_scheme = _make_scheme(context, target_boundary, stream=stream) + input_ordering = _make_ordering(context, 4, stream=stream) + output_ordering = _make_ordering(context, target_boundary, stream=stream) output_by_pid = asyncio.run( _adjust_and_collect( context, comm, _frame(list(range(8))), - input_scheme, - output_scheme, + input_ordering, + output_ordering, ) ) @@ -449,22 +444,22 @@ def test_adjust_orderscheme_single_rank_no_collective( @pytest.mark.spmd -def test_adjust_orderscheme_multi_chunk_input(spmd_engine: SPMDEngine) -> None: +def test_adjust_ordering_multi_chunk_input(spmd_engine: SPMDEngine) -> None: context = spmd_engine.context comm = spmd_engine.comm if comm.nranks != 1: pytest.skip("This test covers local chunk accumulation.") stream = context.br().stream_pool.get_stream() - input_scheme = _make_scheme(context, 4, stream=stream) - output_scheme = _make_scheme(context, 4, stream=stream) + input_ordering = _make_ordering(context, 4, stream=stream) + output_ordering = _make_ordering(context, 4, stream=stream) output = asyncio.run( _adjust_and_collect( context, comm, [_frame([0, 1]), _frame([2, 3, 4, 5]), _frame([6, 7])], - input_scheme, - output_scheme, + input_ordering, + output_ordering, ) ) From 11e1501ddf3c1a8737145d83225cd866036e49f3 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 9 Jul 2026 08:09:44 -0700 Subject: [PATCH 14/26] cleanup --- .../actor_graph/collectives/ordering.py | 2 ++ .../tests/streaming/test_adjust_ordering.py | 32 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 23a643350891..8c18af3e60de 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -251,6 +251,8 @@ def _has_reached(self, stop: int) -> bool: async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: """Return all locally-read pieces for output pids in ``[start, stop)``.""" + if start >= stop: + return {} while not self.input_done and not self._has_reached(stop): msg = await self.ch_in.recv(self.context) if msg is None: diff --git a/python/cudf_polars/tests/streaming/test_adjust_ordering.py b/python/cudf_polars/tests/streaming/test_adjust_ordering.py index d7746816a215..86158044163b 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_ordering.py +++ b/python/cudf_polars/tests/streaming/test_adjust_ordering.py @@ -369,6 +369,38 @@ def test_adjust_ordering_middle_rank_buffers_only_as_needed( _assert_partition_output(output, expected) +@pytest.mark.spmd +def test_adjust_ordering_empty_rank_window(spmd_engine: SPMDEngine) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks != 3: + pytest.skip("This test expects exactly three ranks.") + + keys = {0: [0, 1], 1: [5, 6], 2: [9, 10]}[comm.rank] + stream = context.br().stream_pool.get_stream() + input_ordering = _make_ordering(context, [5, 9], stream=stream) + output_ordering = _make_ordering(context, 5, stream=stream) + + with reserve_op_id() as op_id: + output = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(keys), + input_ordering, + output_ordering, + collective_id=op_id, + ) + ) + + expected = { + 0: {0: [0, 1]}, + 1: {1: [5, 6, 9, 10]}, + 2: {}, + }[comm.rank] + _assert_partition_output(output, expected) + + @pytest.mark.spmd def test_adjust_ordering_all_empty_input(spmd_engine: SPMDEngine) -> None: context = spmd_engine.context From 54c8b02a0d8515a16402e18d2524f1f2f856d748 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 9 Jul 2026 12:26:01 -0700 Subject: [PATCH 15/26] first thorough cleanup pass --- .../actor_graph/collectives/ordering.py | 312 +++++++++++------- 1 file changed, 190 insertions(+), 122 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 8c18af3e60de..111b2c66ccd3 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Ordering adjustment utilities for the RapidsMPF streaming runtime.""" +"""Adjust streams between concrete Ordering boundary layouts without sorting.""" from __future__ import annotations @@ -57,9 +57,9 @@ def _partition_range(rank: int, nranks: int, npartitions: int) -> tuple[int, int def _validate_orderings(input_ordering: Ordering, output_ordering: Ordering) -> None: - """Validate the first-pass flat Ordering adjustment contract.""" + """Validate the Ordering pair supported by this data-movement primitive.""" if not output_ordering.strict_boundaries: - raise ValueError("adjust_ordering requires a strict output Ordering.") + raise ValueError("adjust_ordering requires strict output boundaries.") prefix_len = len(output_ordering.keys) if input_ordering.keys[:prefix_len] != output_ordering.keys: raise NotImplementedError( @@ -98,7 +98,7 @@ def _split_points( def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Table: - """Append a hidden target-partition-id column to *table*.""" + """Append a temporary target-partition-id column to *table*.""" pid_col = plc.Column.from_scalar( plc.Scalar.from_py(pid, _PID_PLC_DTYPE, stream=stream), table.num_rows(), @@ -120,6 +120,7 @@ def _boundary_search_positions( input_prefix_boundaries = plc.Table(input_boundary_table.columns()[:prefix_len]) orders = [key.order for key in output_ordering.keys] null_orders = [key.null_order for key in output_ordering.keys] + # Keep both sides of equal-boundary runs for prefix/non-strict cases. lower_col = plc.search.lower_bound( output_boundary_table, input_prefix_boundaries, @@ -155,6 +156,7 @@ def _source_output_range( input_npartitions = input_ordering.num_boundaries + 1 output_npartitions = output_ordering.num_boundaries + 1 output_prefix_only = len(output_ordering.keys) < len(input_ordering.keys) + # Prefix/non-strict boundaries can overlap the equal-boundary run above them. include_upper_boundary = output_prefix_only or not input_ordering.strict_boundaries input_start, input_stop = _partition_range(source_rank, nranks, input_npartitions) if input_start == input_stop: @@ -182,7 +184,7 @@ def _unpack_remote_piece( stream: Stream, br: BufferResource, ) -> tuple[int, TableChunk] | None: - """Unpack one remote piece and recover its hidden target partition ID.""" + """Unpack one remote piece and recover its temporary target partition ID.""" table = unpack_and_concat([packed], stream=stream, br=br) if table.num_rows() == 0: return None @@ -197,9 +199,7 @@ def _unpack_remote_piece( .to_polars() .item(0, 0) ) - payload = plc.concatenate.concatenate( - [plc.Table(payload_cols)], stream=stream, mr=br.device_mr - ) + payload = plc.Table(payload_cols).copy(stream=stream, mr=br.device_mr) return pid, TableChunk.from_pylibcudf_table( payload, stream, @@ -214,7 +214,7 @@ def _copy_to_owned_chunk( br: BufferResource, ) -> TableChunk: """Copy a table view into a uniquely-owned chunk.""" - table = plc.concatenate.concatenate([table], stream=stream, mr=br.device_mr) + table = table.copy(stream=stream, mr=br.device_mr) return TableChunk.from_pylibcudf_table( table, stream, @@ -236,24 +236,15 @@ def __init__( self.context = context self.ch_in = ch_in self.boundary_chunk = boundary_chunk - self.boundary_table = boundary_chunk.table_view() self.output_ordering = output_ordering self.pending: dict[int, ChunkStore] = {} self.input_done = False - def _store(self, pid: int, chunk: TableChunk) -> None: - if pid not in self.pending: - self.pending[pid] = ChunkStore(self.context) - self.pending[pid].insert(Message(pid, chunk)) - - def _has_reached(self, stop: int) -> bool: - return any(pid >= stop for pid in self.pending) - async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: """Return all locally-read pieces for output pids in ``[start, stop)``.""" if start >= stop: return {} - while not self.input_done and not self._has_reached(stop): + while not self.input_done and not any(pid >= stop for pid in self.pending): msg = await self.ch_in.recv(self.context) if msg is None: self.input_done = True @@ -269,15 +260,21 @@ async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: ) as stream: table = chunk.table_view() splits = _split_points( - table, self.boundary_table, self.output_ordering, stream + table, + self.boundary_chunk.table_view(), + self.output_ordering, + stream, ) for pid, piece in enumerate( plc.copying.split(table, splits, stream=stream) ): if piece.num_rows() == 0: continue - self._store( - pid, _copy_to_owned_chunk(piece, stream, self.context.br()) + _store_chunk( + self.context, + self.pending, + pid, + _copy_to_owned_chunk(piece, stream, self.context.br()), ) out = { @@ -296,6 +293,7 @@ async def _adjust_ordering_local( ch_in: Channel[TableChunk], output_ordering: Ordering, ) -> None: + """Adjust ordering on one rank without draining all input before emitting.""" npartitions = output_ordering.num_boundaries + 1 boundary_chunk = output_ordering.get_boundaries(context.br()) boundary_table = boundary_chunk.table_view() @@ -368,7 +366,145 @@ def _store_chunk( stores[pid].insert(Message(pid, chunk)) -async def _adjust_ordering_rank_hybrid( +async def _send_remote_piece( + context: Context, + comm: Communicator, + exchange: SparseAlltoall, + npartitions: int, + pid: int, + chunk: TableChunk, +) -> None: + """Send one output-partition piece to its remote owner.""" + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(chunk.stream,), + ) as stream: + exchange.insert( + _contiguous_owner(pid, comm.nranks, npartitions), + packed_data_from_cudf_packed_columns( + pack( + _append_partition_id(chunk.table_view(), pid, stream), + stream, + mr=context.br().device_mr, + ), + stream, + context.br(), + ), + ) + + +async def _emit_partition( + context: Context, + ref_ir: IR, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + pid: int, + store: ChunkStore | None, +) -> None: + """Emit one output partition, using an empty chunk when no data is present.""" + chunks = ( + [TableChunk.from_message(msg, br=context.br()) for msg in store] + if store is not None + else [] + ) + chunk = ( + await concat_batch(chunks, context, ref_ir.schema, ir_context) + if chunks + else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) + ) + await ch_out.send(context, Message(pid, chunk)) + + +async def _adjust_ordering_unblocked_rank( + context: Context, + comm: Communicator, + ref_ir: IR, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + output_ordering: Ordering, + exchange: SparseAlltoall, + local_window: tuple[int, int], + npartitions: int, + boundary_chunk: TableChunk, +) -> None: + """Stream a rank that does not need remote input before emitting.""" + pending_pid: int | None = None + pending_chunks: ChunkStore | None = None + next_pid = local_window[0] + + async def emit_pending(pid: int) -> None: + nonlocal pending_pid, pending_chunks + await _emit_partition( + context, + ref_ir, + ir_context, + ch_out, + pid, + pending_chunks if pending_pid == pid else None, + ) + if pending_pid == pid: + pending_pid = None + pending_chunks = None + + try: + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + if chunk.table_view().num_rows() == 0: + continue + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(chunk.stream, boundary_chunk.stream), + ) as stream: + table = chunk.table_view() + splits = _split_points( + table, + boundary_chunk.table_view(), + output_ordering, + stream, + ) + for pid, piece in enumerate( + plc.copying.split(table, splits, stream=stream) + ): + if piece.num_rows() == 0: + continue + owner = _contiguous_owner(pid, comm.nranks, npartitions) + owned = owner == comm.rank + piece_chunk = _copy_to_owned_chunk(piece, stream, context.br()) + if not owned: + await _send_remote_piece( + context, + comm, + exchange, + npartitions, + pid, + piece_chunk, + ) + continue + if pending_pid is not None and pending_pid != pid: + emitted_pid = pending_pid + await emit_pending(emitted_pid) + next_pid = emitted_pid + 1 + while next_pid < pid: + await emit_pending(next_pid) + next_pid += 1 + if pending_pid is None: + pending_pid = pid + pending_chunks = ChunkStore(context) + assert pending_chunks is not None + pending_chunks.insert(Message(pid, piece_chunk)) + finally: + await exchange.insert_finished(context) + + while next_pid < local_window[1]: + await emit_pending(next_pid) + next_pid += 1 + await ch_out.drain(context) + + +async def _adjust_ordering_multi_rank( context: Context, comm: Communicator, ref_ir: IR, @@ -381,9 +517,9 @@ async def _adjust_ordering_rank_hybrid( lower_positions: list[int], upper_positions: list[int], ) -> None: + """Adjust ordering across ranks while buffering only blocked ranks.""" npartitions = output_ordering.num_boundaries + 1 boundary_chunk = output_ordering.get_boundaries(context.br()) - boundary_table = boundary_chunk.table_view() local_window = _partition_range(comm.rank, comm.nranks, npartitions) source_ranges = [ _source_output_range( @@ -397,12 +533,12 @@ async def _adjust_ordering_rank_hybrid( for source_rank in range(comm.nranks) ] local_source_range = source_ranges[comm.rank] - srcs = [ + remote_sources = [ source_rank for source_rank, source_range in enumerate(source_ranges) if source_rank != comm.rank and _ranges_overlap(source_range, local_window) ] - dsts = [ + remote_destinations = [ output_rank for output_rank in range(comm.nranks) if output_rank != comm.rank @@ -414,107 +550,33 @@ async def _adjust_ordering_rank_hybrid( context, comm, collective_id, - srcs=srcs, - dsts=dsts, + srcs=remote_sources, + dsts=remote_destinations, ) - async def send_piece(pid: int, chunk: TableChunk) -> None: - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(chunk.stream,), - ) as stream: - exchange.insert( - _contiguous_owner(pid, comm.nranks, npartitions), - packed_data_from_cudf_packed_columns( - pack( - _append_partition_id(chunk.table_view(), pid, stream), - stream, - mr=context.br().device_mr, - ), - stream, - context.br(), - ), - ) - - async def emit(pid: int, store: ChunkStore | None) -> None: - chunks = ( - [TableChunk.from_message(msg, br=context.br()) for msg in store] - if store is not None - else [] - ) - chunk = ( - await concat_batch(chunks, context, ref_ir.schema, ir_context) - if chunks - else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) - ) - await ch_out.send(context, Message(pid, chunk)) - # Ranks with no incoming dependency can stream local output immediately # while sending remote-owned pieces as they are encountered. - if not srcs: - pending_pid: int | None = None - pending_chunks: ChunkStore | None = None - next_pid = local_window[0] - - async def emit_pending(pid: int) -> None: - nonlocal pending_pid, pending_chunks - await emit(pid, pending_chunks if pending_pid == pid else None) - if pending_pid == pid: - pending_pid = None - pending_chunks = None - - try: - while (msg := await ch_in.recv(context)) is not None: - chunk = TableChunk.from_message( - msg, br=context.br() - ).make_available_and_spill(context.br(), allow_overbooking=True) - if chunk.table_view().num_rows() == 0: - continue - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(chunk.stream, boundary_chunk.stream), - ) as stream: - table = chunk.table_view() - splits = _split_points( - table, boundary_table, output_ordering, stream - ) - for pid, piece in enumerate( - plc.copying.split(table, splits, stream=stream) - ): - if piece.num_rows() == 0: - continue - owner = _contiguous_owner(pid, comm.nranks, npartitions) - owned = owner == comm.rank - piece_chunk = _copy_to_owned_chunk(piece, stream, context.br()) - if not owned: - await send_piece(pid, piece_chunk) - continue - if pending_pid is not None and pending_pid != pid: - emitted_pid = pending_pid - await emit_pending(emitted_pid) - next_pid = emitted_pid + 1 - while next_pid < pid: - await emit_pending(next_pid) - next_pid += 1 - if pending_pid is None: - pending_pid = pid - pending_chunks = ChunkStore(context) - assert pending_chunks is not None - pending_chunks.insert(Message(pid, piece_chunk)) - finally: - await exchange.insert_finished(context) - - while next_pid < local_window[1]: - await emit_pending(next_pid) - next_pid += 1 - await ch_out.drain(context) + if not remote_sources: + await _adjust_ordering_unblocked_rank( + context, + comm, + ref_ir, + ir_context, + ch_out, + ch_in, + output_ordering, + exchange, + local_window, + npartitions, + boundary_chunk, + ) return reader = _OutputPieceReader(context, ch_in, boundary_chunk, output_ordering) local_pieces: dict[int, ChunkStore] = {} # If a higher rank depends on this rank's data, read far enough to make # that data available before waiting for lower-rank input. - if dsts: + if remote_destinations: pieces = await reader.collect_window(*local_source_range) for pid, store in pieces.items(): owner = _contiguous_owner(pid, comm.nranks, npartitions) @@ -522,13 +584,20 @@ async def emit_pending(pid: int) -> None: local_pieces[pid] = store continue for msg in store: - await send_piece(pid, TableChunk.from_message(msg, br=context.br())) + await _send_remote_piece( + context, + comm, + exchange, + npartitions, + pid, + TableChunk.from_message(msg, br=context.br()), + ) await exchange.insert_finished(context) pieces_by_source: dict[int, dict[int, ChunkStore]] = {} if local_pieces: pieces_by_source[comm.rank] = local_pieces - for source_rank in srcs: + for source_rank in remote_sources: remote_pieces: dict[int, ChunkStore] = {} stream = context.br().stream_pool.get_stream() for packed in exchange.extract(source_rank): @@ -636,19 +705,18 @@ async def adjust_ordering( input_boundary_chunk = input_ordering.get_boundaries(context.br()) boundary_chunk = output_ordering.get_boundaries(context.br()) - boundary_table = boundary_chunk.table_view() with stream_ordered_after( context.br().stream_pool.get_stream, upstreams=(input_boundary_chunk.stream, boundary_chunk.stream), ) as stream: lower_positions, upper_positions = _boundary_search_positions( input_boundary_chunk.table_view(), - boundary_table, + boundary_chunk.table_view(), output_ordering, stream, ) assert collective_id is not None - await _adjust_ordering_rank_hybrid( + await _adjust_ordering_multi_rank( context, comm, ref_ir, From a8c38def3bbc12ee604eec8e0cd5f731ea3b65cc Mon Sep 17 00:00:00 2001 From: rjzamora Date: Thu, 9 Jul 2026 12:55:05 -0700 Subject: [PATCH 16/26] unify single- and multi-rnk --- .../actor_graph/collectives/ordering.py | 332 ++++++++---------- 1 file changed, 141 insertions(+), 191 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 111b2c66ccd3..2747c4ee28d9 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -28,6 +28,8 @@ from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + from cudf_streaming.channel_metadata import Ordering from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.memory.buffer_resource import BufferResource @@ -285,71 +287,92 @@ async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: return dict(sorted(out.items())) -async def _adjust_ordering_local( +async def _adjust_ordering_streaming_window( context: Context, ref_ir: IR, ir_context: IRExecutionContext, ch_out: Channel[TableChunk], ch_in: Channel[TableChunk], output_ordering: Ordering, + local_window: tuple[int, int], + boundary_chunk: TableChunk, + owns_pid: Callable[[int], bool], + send_remote_piece: Callable[[int, TableChunk], Awaitable[None]] | None = None, + finish_remote_sends: Callable[[], Awaitable[None]] | None = None, ) -> None: - """Adjust ordering on one rank without draining all input before emitting.""" - npartitions = output_ordering.num_boundaries + 1 - boundary_chunk = output_ordering.get_boundaries(context.br()) - boundary_table = boundary_chunk.table_view() + """Stream one output window without draining all input before emitting.""" pending_pid: int | None = None pending_chunks: ChunkStore | None = None - next_pid = 0 + next_pid = local_window[0] async def emit_pending(pid: int) -> None: nonlocal pending_pid, pending_chunks - if pending_pid == pid and pending_chunks is not None: - chunks = [ - TableChunk.from_message(msg, br=context.br()) for msg in pending_chunks - ] - chunk = await concat_batch(chunks, context, ref_ir.schema, ir_context) + await _emit_partition( + context, + ref_ir, + ir_context, + ch_out, + pid, + pending_chunks if pending_pid == pid else None, + ) + if pending_pid == pid: pending_pid = None pending_chunks = None - else: - chunk = empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) - await ch_out.send(context, Message(pid, chunk)) - while (msg := await ch_in.recv(context)) is not None: - chunk = TableChunk.from_message(msg, br=context.br()).make_available_and_spill( - context.br(), allow_overbooking=True - ) - if chunk.table_view().num_rows() == 0: - continue - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(chunk.stream, boundary_chunk.stream), - ) as stream: - table = chunk.table_view() - splits = _split_points(table, boundary_table, output_ordering, stream) - for pid, piece in enumerate( - plc.copying.split(table, splits, stream=stream) - ): - if piece.num_rows() == 0: - continue - if pending_pid is not None and pending_pid != pid: - emitted_pid = pending_pid - await emit_pending(emitted_pid) - next_pid = emitted_pid + 1 - while next_pid < pid: - await emit_pending(next_pid) - next_pid += 1 - if pending_pid is None: - pending_pid = pid - pending_chunks = ChunkStore(context) - assert pending_chunks is not None - pending_chunks.insert( - Message( - pid, - _copy_to_owned_chunk(piece, stream, context.br()), - ) + try: + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + if chunk.table_view().num_rows() == 0: + continue + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(chunk.stream, boundary_chunk.stream), + ) as stream: + table = chunk.table_view() + splits = _split_points( + table, + boundary_chunk.table_view(), + output_ordering, + stream, ) + for pid, piece in enumerate( + plc.copying.split(table, splits, stream=stream) + ): + if piece.num_rows() == 0: + continue + piece_chunk = _copy_to_owned_chunk(piece, stream, context.br()) + if not owns_pid(pid): + if send_remote_piece is None: + raise RuntimeError( + "Encountered remote-owned ordering piece without " + "a remote sender." + ) + await send_remote_piece(pid, piece_chunk) + continue + if pending_pid is not None and pending_pid != pid: + emitted_pid = pending_pid + await emit_pending(emitted_pid) + next_pid = emitted_pid + 1 + while next_pid < pid: + await emit_pending(next_pid) + next_pid += 1 + if pending_pid is None: + pending_pid = pid + pending_chunks = ChunkStore(context) + assert pending_chunks is not None + pending_chunks.insert( + Message( + pid, + piece_chunk, + ) + ) + finally: + if finish_remote_sends is not None: + await finish_remote_sends() - while next_pid < npartitions: + while next_pid < local_window[1]: await emit_pending(next_pid) next_pid += 1 await ch_out.drain(context) @@ -415,96 +438,7 @@ async def _emit_partition( await ch_out.send(context, Message(pid, chunk)) -async def _adjust_ordering_unblocked_rank( - context: Context, - comm: Communicator, - ref_ir: IR, - ir_context: IRExecutionContext, - ch_out: Channel[TableChunk], - ch_in: Channel[TableChunk], - output_ordering: Ordering, - exchange: SparseAlltoall, - local_window: tuple[int, int], - npartitions: int, - boundary_chunk: TableChunk, -) -> None: - """Stream a rank that does not need remote input before emitting.""" - pending_pid: int | None = None - pending_chunks: ChunkStore | None = None - next_pid = local_window[0] - - async def emit_pending(pid: int) -> None: - nonlocal pending_pid, pending_chunks - await _emit_partition( - context, - ref_ir, - ir_context, - ch_out, - pid, - pending_chunks if pending_pid == pid else None, - ) - if pending_pid == pid: - pending_pid = None - pending_chunks = None - - try: - while (msg := await ch_in.recv(context)) is not None: - chunk = TableChunk.from_message( - msg, br=context.br() - ).make_available_and_spill(context.br(), allow_overbooking=True) - if chunk.table_view().num_rows() == 0: - continue - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(chunk.stream, boundary_chunk.stream), - ) as stream: - table = chunk.table_view() - splits = _split_points( - table, - boundary_chunk.table_view(), - output_ordering, - stream, - ) - for pid, piece in enumerate( - plc.copying.split(table, splits, stream=stream) - ): - if piece.num_rows() == 0: - continue - owner = _contiguous_owner(pid, comm.nranks, npartitions) - owned = owner == comm.rank - piece_chunk = _copy_to_owned_chunk(piece, stream, context.br()) - if not owned: - await _send_remote_piece( - context, - comm, - exchange, - npartitions, - pid, - piece_chunk, - ) - continue - if pending_pid is not None and pending_pid != pid: - emitted_pid = pending_pid - await emit_pending(emitted_pid) - next_pid = emitted_pid + 1 - while next_pid < pid: - await emit_pending(next_pid) - next_pid += 1 - if pending_pid is None: - pending_pid = pid - pending_chunks = ChunkStore(context) - assert pending_chunks is not None - pending_chunks.insert(Message(pid, piece_chunk)) - finally: - await exchange.insert_finished(context) - - while next_pid < local_window[1]: - await emit_pending(next_pid) - next_pid += 1 - await ch_out.drain(context) - - -async def _adjust_ordering_multi_rank( +async def _adjust_ordering_impl( context: Context, comm: Communicator, ref_ir: IR, @@ -513,25 +447,38 @@ async def _adjust_ordering_multi_rank( ch_in: Channel[TableChunk], input_ordering: Ordering, output_ordering: Ordering, - collective_id: int, - lower_positions: list[int], - upper_positions: list[int], + collective_id: int | None, ) -> None: - """Adjust ordering across ranks while buffering only blocked ranks.""" + """Adjust ordering while using exchange only for remote dependencies.""" npartitions = output_ordering.num_boundaries + 1 boundary_chunk = output_ordering.get_boundaries(context.br()) local_window = _partition_range(comm.rank, comm.nranks, npartitions) - source_ranges = [ - _source_output_range( - source_rank, - comm.nranks, - input_ordering, - output_ordering, - lower_positions, - upper_positions, - ) - for source_rank in range(comm.nranks) - ] + + if comm.nranks == 1: + source_ranges = [(0, npartitions)] + else: + input_boundary_chunk = input_ordering.get_boundaries(context.br()) + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(input_boundary_chunk.stream, boundary_chunk.stream), + ) as stream: + lower_positions, upper_positions = _boundary_search_positions( + input_boundary_chunk.table_view(), + boundary_chunk.table_view(), + output_ordering, + stream, + ) + source_ranges = [ + _source_output_range( + source_rank, + comm.nranks, + input_ordering, + output_ordering, + lower_positions, + upper_positions, + ) + for source_rank in range(comm.nranks) + ] local_source_range = source_ranges[comm.rank] remote_sources = [ source_rank @@ -546,32 +493,61 @@ async def _adjust_ordering_multi_rank( local_source_range, _partition_range(output_rank, comm.nranks, npartitions) ) ] - exchange = SparseAlltoall( - context, - comm, - collective_id, - srcs=remote_sources, - dsts=remote_destinations, - ) # Ranks with no incoming dependency can stream local output immediately # while sending remote-owned pieces as they are encountered. if not remote_sources: - await _adjust_ordering_unblocked_rank( + if remote_destinations: + assert collective_id is not None + exchange = SparseAlltoall( + context, + comm, + collective_id, + srcs=[], + dsts=remote_destinations, + ) + else: + exchange = None + + async def send_piece(pid: int, chunk: TableChunk) -> None: + assert exchange is not None + await _send_remote_piece( + context, + comm, + exchange, + npartitions, + pid, + chunk, + ) + + async def finish_sends() -> None: + assert exchange is not None + await exchange.insert_finished(context) + + await _adjust_ordering_streaming_window( context, - comm, ref_ir, ir_context, ch_out, ch_in, output_ordering, - exchange, local_window, - npartitions, boundary_chunk, + lambda pid: _contiguous_owner(pid, comm.nranks, npartitions) == comm.rank, + send_piece if exchange is not None else None, + finish_sends if exchange is not None else None, ) return + assert collective_id is not None + exchange = SparseAlltoall( + context, + comm, + collective_id, + srcs=remote_sources, + dsts=remote_destinations, + ) + reader = _OutputPieceReader(context, ch_in, boundary_chunk, output_ordering) local_pieces: dict[int, ChunkStore] = {} # If a higher rank depends on this rank's data, read far enough to make @@ -692,31 +668,7 @@ async def adjust_ordering( raise ValueError("collective_id is required when comm.nranks > 1.") try: - if comm.nranks == 1: - await _adjust_ordering_local( - context, - ref_ir, - ir_context, - ch_out, - ch_in, - output_ordering, - ) - return - - input_boundary_chunk = input_ordering.get_boundaries(context.br()) - boundary_chunk = output_ordering.get_boundaries(context.br()) - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(input_boundary_chunk.stream, boundary_chunk.stream), - ) as stream: - lower_positions, upper_positions = _boundary_search_positions( - input_boundary_chunk.table_view(), - boundary_chunk.table_view(), - output_ordering, - stream, - ) - assert collective_id is not None - await _adjust_ordering_multi_rank( + await _adjust_ordering_impl( context, comm, ref_ir, @@ -726,8 +678,6 @@ async def adjust_ordering( input_ordering, output_ordering, collective_id, - lower_positions, - upper_positions, ) except BaseException: await gather_in_task_group( From fd1fd3caf0bb28619031cfbba104fc1cf8b40bb5 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 10 Jul 2026 17:50:28 -0700 Subject: [PATCH 17/26] save new experimental design --- .../actor_graph/collectives/ordering.py | 151 ++++++++++++++---- 1 file changed, 122 insertions(+), 29 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 2747c4ee28d9..dacee15fddab 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -181,6 +181,18 @@ def _ranges_overlap(left: tuple[int, int], right: tuple[int, int]) -> bool: return max(left[0], right[0]) < min(left[1], right[1]) +def _sources_for_pid( + source_ranges: list[tuple[int, int]], + pid: int, +) -> list[int]: + """Return source ranks that may contribute to one output partition.""" + return [ + source_rank + for source_rank, source_range in enumerate(source_ranges) + if source_range[0] <= pid < source_range[1] + ] + + def _unpack_remote_piece( packed: PackedData, stream: Stream, @@ -201,6 +213,8 @@ def _unpack_remote_piece( .to_polars() .item(0, 0) ) + if not payload_cols: + return None payload = plc.Table(payload_cols).copy(stream=stream, mr=br.device_mr) return pid, TableChunk.from_pylibcudf_table( payload, @@ -286,6 +300,12 @@ async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: } return dict(sorted(out.items())) + async def drain_remaining(self) -> None: + """Consume input that cannot affect any remaining output partition.""" + while not self.input_done: + if await self.ch_in.recv(self.context) is None: + self.input_done = True + async def _adjust_ordering_streaming_window( context: Context, @@ -416,6 +436,33 @@ async def _send_remote_piece( ) +async def _send_remote_marker( + context: Context, + comm: Communicator, + exchange: SparseAlltoall, + npartitions: int, + pid: int, +) -> None: + """Send a no-payload marker for an empty remote-owned partition.""" + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(), + ) as stream: + pid_col = plc.Column.from_scalar( + plc.Scalar.from_py(pid, _PID_PLC_DTYPE, stream=stream), + 1, + stream=stream, + ) + exchange.insert( + _contiguous_owner(pid, comm.nranks, npartitions), + packed_data_from_cudf_packed_columns( + pack(plc.Table([pid_col]), stream, mr=context.br().device_mr), + stream, + context.br(), + ), + ) + + async def _emit_partition( context: Context, ref_ir: IR, @@ -485,14 +532,17 @@ async def _adjust_ordering_impl( for source_rank, source_range in enumerate(source_ranges) if source_rank != comm.rank and _ranges_overlap(source_range, local_window) ] - remote_destinations = [ - output_rank - for output_rank in range(comm.nranks) - if output_rank != comm.rank - and _ranges_overlap( - local_source_range, _partition_range(output_rank, comm.nranks, npartitions) - ) + owed_remote_pids = [ + pid + for pid in range(*local_source_range) + if _contiguous_owner(pid, comm.nranks, npartitions) != comm.rank ] + remote_destinations = sorted( + {_contiguous_owner(pid, comm.nranks, npartitions) for pid in owed_remote_pids} + ) + owed_remote_range = ( + (owed_remote_pids[0], owed_remote_pids[-1] + 1) if owed_remote_pids else (0, 0) + ) # Ranks with no incoming dependency can stream local output immediately # while sending remote-owned pieces as they are encountered. @@ -509,8 +559,11 @@ async def _adjust_ordering_impl( else: exchange = None + streamed_remote_pids: set[int] = set() + async def send_piece(pid: int, chunk: TableChunk) -> None: assert exchange is not None + streamed_remote_pids.add(pid) await _send_remote_piece( context, comm, @@ -522,6 +575,15 @@ async def send_piece(pid: int, chunk: TableChunk) -> None: async def finish_sends() -> None: assert exchange is not None + for pid in owed_remote_pids: + if pid not in streamed_remote_pids: + await _send_remote_marker( + context, + comm, + exchange, + npartitions, + pid, + ) await exchange.insert_finished(context) await _adjust_ordering_streaming_window( @@ -550,16 +612,40 @@ async def finish_sends() -> None: reader = _OutputPieceReader(context, ch_in, boundary_chunk, output_ordering) local_pieces: dict[int, ChunkStore] = {} - # If a higher rank depends on this rank's data, read far enough to make - # that data available before waiting for lower-rank input. - if remote_destinations: - pieces = await reader.collect_window(*local_source_range) + buffered_remote_pids: set[int] = set() + first_blocked_pid = next( + ( + pid + for pid in range(*local_window) + if any( + source_rank != comm.rank + for source_rank in _sources_for_pid(source_ranges, pid) + ) + ), + local_window[1], + ) + if first_blocked_pid > local_window[0]: + prefix_pieces = await reader.collect_window(local_window[0], first_blocked_pid) + for pid in range(local_window[0], first_blocked_pid): + await _emit_partition( + context, + ref_ir, + ir_context, + ch_out, + pid, + prefix_pieces.get(pid), + ) + + # Read only far enough to satisfy destination liabilities before receiving. + if owed_remote_pids: + pieces = await reader.collect_window(*owed_remote_range) for pid, store in pieces.items(): owner = _contiguous_owner(pid, comm.nranks, npartitions) if owner == comm.rank: local_pieces[pid] = store continue for msg in store: + buffered_remote_pids.add(pid) await _send_remote_piece( context, comm, @@ -568,11 +654,19 @@ async def finish_sends() -> None: pid, TableChunk.from_message(msg, br=context.br()), ) + for pid in owed_remote_pids: + if pid not in buffered_remote_pids: + await _send_remote_marker( + context, + comm, + exchange, + npartitions, + pid, + ) await exchange.insert_finished(context) pieces_by_source: dict[int, dict[int, ChunkStore]] = {} - if local_pieces: - pieces_by_source[comm.rank] = local_pieces + pieces_by_source[comm.rank] = local_pieces for source_rank in remote_sources: remote_pieces: dict[int, ChunkStore] = {} stream = context.br().stream_pool.get_stream() @@ -584,23 +678,21 @@ async def finish_sends() -> None: _store_chunk(context, remote_pieces, pid, chunk) pieces_by_source[source_rank] = remote_pieces - for pid, store in (await reader.collect_window(*local_window)).items(): - if _contiguous_owner(pid, comm.nranks, npartitions) == comm.rank: - if pid not in local_pieces: - local_pieces[pid] = ChunkStore(context) - for msg in store: - local_pieces[pid].insert(msg) - if local_pieces: - pieces_by_source[comm.rank] = local_pieces - - contributing_sources = [ - source_rank - for source_rank, source_range in enumerate(source_ranges) - if _ranges_overlap(source_range, local_window) - ] - for pid in range(*local_window): + for pid in range(first_blocked_pid, local_window[1]): + pid_sources = _sources_for_pid(source_ranges, pid) + if comm.rank in pid_sources and pid not in local_pieces: + local_pieces.update( + { + piece_pid: store + for piece_pid, store in ( + await reader.collect_window(pid, pid + 1) + ).items() + if _contiguous_owner(piece_pid, comm.nranks, npartitions) + == comm.rank + } + ) chunks: list[TableChunk] = [] - for source_rank in contributing_sources: + for source_rank in pid_sources: stores = pieces_by_source.get(source_rank) if stores is None: continue @@ -616,6 +708,7 @@ async def finish_sends() -> None: else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) ) await ch_out.send(context, Message(pid, chunk)) + await reader.drain_remaining() await ch_out.drain(context) From fd34b383d736d84f915813a87300ad0c8c489694 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Fri, 10 Jul 2026 18:00:56 -0700 Subject: [PATCH 18/26] try cleaning up - impl is still many lines --- .../actor_graph/collectives/ordering.py | 313 ++++++++---------- 1 file changed, 145 insertions(+), 168 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index dacee15fddab..f1349ec0d061 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -4,6 +4,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING import polars as pl @@ -28,8 +29,6 @@ from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: - from collections.abc import Awaitable, Callable - from cudf_streaming.channel_metadata import Ordering from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.memory.buffer_resource import BufferResource @@ -45,6 +44,18 @@ _PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32) +@dataclass(frozen=True) +class _RoutingPlan: + npartitions: int + boundary_chunk: TableChunk + local_window: tuple[int, int] + source_ranges: list[tuple[int, int]] + remote_sources: list[int] + owed_remote_pids: list[int] + remote_destinations: list[int] + owed_remote_range: tuple[int, int] + + def _contiguous_owner(pid: int, nranks: int, npartitions: int) -> int: """Return the rank owning *pid* under contiguous partition assignment.""" return pid * nranks // npartitions @@ -307,97 +318,6 @@ async def drain_remaining(self) -> None: self.input_done = True -async def _adjust_ordering_streaming_window( - context: Context, - ref_ir: IR, - ir_context: IRExecutionContext, - ch_out: Channel[TableChunk], - ch_in: Channel[TableChunk], - output_ordering: Ordering, - local_window: tuple[int, int], - boundary_chunk: TableChunk, - owns_pid: Callable[[int], bool], - send_remote_piece: Callable[[int, TableChunk], Awaitable[None]] | None = None, - finish_remote_sends: Callable[[], Awaitable[None]] | None = None, -) -> None: - """Stream one output window without draining all input before emitting.""" - pending_pid: int | None = None - pending_chunks: ChunkStore | None = None - next_pid = local_window[0] - - async def emit_pending(pid: int) -> None: - nonlocal pending_pid, pending_chunks - await _emit_partition( - context, - ref_ir, - ir_context, - ch_out, - pid, - pending_chunks if pending_pid == pid else None, - ) - if pending_pid == pid: - pending_pid = None - pending_chunks = None - - try: - while (msg := await ch_in.recv(context)) is not None: - chunk = TableChunk.from_message( - msg, br=context.br() - ).make_available_and_spill(context.br(), allow_overbooking=True) - if chunk.table_view().num_rows() == 0: - continue - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(chunk.stream, boundary_chunk.stream), - ) as stream: - table = chunk.table_view() - splits = _split_points( - table, - boundary_chunk.table_view(), - output_ordering, - stream, - ) - for pid, piece in enumerate( - plc.copying.split(table, splits, stream=stream) - ): - if piece.num_rows() == 0: - continue - piece_chunk = _copy_to_owned_chunk(piece, stream, context.br()) - if not owns_pid(pid): - if send_remote_piece is None: - raise RuntimeError( - "Encountered remote-owned ordering piece without " - "a remote sender." - ) - await send_remote_piece(pid, piece_chunk) - continue - if pending_pid is not None and pending_pid != pid: - emitted_pid = pending_pid - await emit_pending(emitted_pid) - next_pid = emitted_pid + 1 - while next_pid < pid: - await emit_pending(next_pid) - next_pid += 1 - if pending_pid is None: - pending_pid = pid - pending_chunks = ChunkStore(context) - assert pending_chunks is not None - pending_chunks.insert( - Message( - pid, - piece_chunk, - ) - ) - finally: - if finish_remote_sends is not None: - await finish_remote_sends() - - while next_pid < local_window[1]: - await emit_pending(next_pid) - next_pid += 1 - await ch_out.drain(context) - - def _store_chunk( context: Context, stores: dict[int, ChunkStore], @@ -409,6 +329,26 @@ def _store_chunk( stores[pid].insert(Message(pid, chunk)) +async def _send_remote_store( + context: Context, + comm: Communicator, + exchange: SparseAlltoall, + npartitions: int, + pid: int, + store: ChunkStore, +) -> None: + """Send all locally-read pieces for one remote-owned output partition.""" + for msg in store: + await _send_remote_piece( + context, + comm, + exchange, + npartitions, + pid, + TableChunk.from_message(msg, br=context.br()), + ) + + async def _send_remote_piece( context: Context, comm: Communicator, @@ -463,6 +403,20 @@ async def _send_remote_marker( ) +async def _send_missing_remote_markers( + context: Context, + comm: Communicator, + exchange: SparseAlltoall, + npartitions: int, + owed_remote_pids: list[int], + sent_remote_pids: set[int], +) -> None: + """Send empty markers for remote-owned partitions with no payload.""" + for pid in owed_remote_pids: + if pid not in sent_remote_pids: + await _send_remote_marker(context, comm, exchange, npartitions, pid) + + async def _emit_partition( context: Context, ref_ir: IR, @@ -485,18 +439,13 @@ async def _emit_partition( await ch_out.send(context, Message(pid, chunk)) -async def _adjust_ordering_impl( +def _make_routing_plan( context: Context, comm: Communicator, - ref_ir: IR, - ir_context: IRExecutionContext, - ch_out: Channel[TableChunk], - ch_in: Channel[TableChunk], input_ordering: Ordering, output_ordering: Ordering, - collective_id: int | None, -) -> None: - """Adjust ordering while using exchange only for remote dependencies.""" +) -> _RoutingPlan: + """Compute local ownership and sparse-exchange obligations.""" npartitions = output_ordering.num_boundaries + 1 boundary_chunk = output_ordering.get_boundaries(context.br()) local_window = _partition_range(comm.rank, comm.nranks, npartitions) @@ -543,62 +492,91 @@ async def _adjust_ordering_impl( owed_remote_range = ( (owed_remote_pids[0], owed_remote_pids[-1] + 1) if owed_remote_pids else (0, 0) ) + return _RoutingPlan( + npartitions=npartitions, + boundary_chunk=boundary_chunk, + local_window=local_window, + source_ranges=source_ranges, + remote_sources=remote_sources, + owed_remote_pids=owed_remote_pids, + remote_destinations=remote_destinations, + owed_remote_range=owed_remote_range, + ) + + +async def _adjust_ordering_impl( + context: Context, + comm: Communicator, + ref_ir: IR, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_in: Channel[TableChunk], + input_ordering: Ordering, + output_ordering: Ordering, + collective_id: int | None, +) -> None: + """Adjust ordering while using exchange only for remote dependencies.""" + plan = _make_routing_plan(context, comm, input_ordering, output_ordering) + reader = _OutputPieceReader(context, ch_in, plan.boundary_chunk, output_ordering) # Ranks with no incoming dependency can stream local output immediately # while sending remote-owned pieces as they are encountered. - if not remote_sources: - if remote_destinations: + if not plan.remote_sources: + if plan.remote_destinations: assert collective_id is not None exchange = SparseAlltoall( context, comm, collective_id, srcs=[], - dsts=remote_destinations, + dsts=plan.remote_destinations, ) else: exchange = None streamed_remote_pids: set[int] = set() - - async def send_piece(pid: int, chunk: TableChunk) -> None: - assert exchange is not None - streamed_remote_pids.add(pid) - await _send_remote_piece( + start, stop = plan.local_window + if plan.owed_remote_pids: + start = min(start, plan.owed_remote_range[0]) + stop = max(stop, plan.owed_remote_range[1]) + for pid in range(start, stop): + pieces = await reader.collect_window(pid, pid + 1) + owner = _contiguous_owner(pid, comm.nranks, plan.npartitions) + if ( + owner == comm.rank + and plan.local_window[0] <= pid < plan.local_window[1] + ): + await _emit_partition( + context, + ref_ir, + ir_context, + ch_out, + pid, + pieces.get(pid), + ) + continue + if exchange is not None and pid in pieces: + streamed_remote_pids.add(pid) + await _send_remote_store( + context, + comm, + exchange, + plan.npartitions, + pid, + pieces[pid], + ) + if exchange is not None: + await _send_missing_remote_markers( context, comm, exchange, - npartitions, - pid, - chunk, + plan.npartitions, + plan.owed_remote_pids, + streamed_remote_pids, ) - - async def finish_sends() -> None: - assert exchange is not None - for pid in owed_remote_pids: - if pid not in streamed_remote_pids: - await _send_remote_marker( - context, - comm, - exchange, - npartitions, - pid, - ) await exchange.insert_finished(context) - - await _adjust_ordering_streaming_window( - context, - ref_ir, - ir_context, - ch_out, - ch_in, - output_ordering, - local_window, - boundary_chunk, - lambda pid: _contiguous_owner(pid, comm.nranks, npartitions) == comm.rank, - send_piece if exchange is not None else None, - finish_sends if exchange is not None else None, - ) + await reader.drain_remaining() + await ch_out.drain(context) return assert collective_id is not None @@ -606,27 +584,28 @@ async def finish_sends() -> None: context, comm, collective_id, - srcs=remote_sources, - dsts=remote_destinations, + srcs=plan.remote_sources, + dsts=plan.remote_destinations, ) - reader = _OutputPieceReader(context, ch_in, boundary_chunk, output_ordering) local_pieces: dict[int, ChunkStore] = {} buffered_remote_pids: set[int] = set() first_blocked_pid = next( ( pid - for pid in range(*local_window) + for pid in range(*plan.local_window) if any( source_rank != comm.rank - for source_rank in _sources_for_pid(source_ranges, pid) + for source_rank in _sources_for_pid(plan.source_ranges, pid) ) ), - local_window[1], + plan.local_window[1], ) - if first_blocked_pid > local_window[0]: - prefix_pieces = await reader.collect_window(local_window[0], first_blocked_pid) - for pid in range(local_window[0], first_blocked_pid): + if first_blocked_pid > plan.local_window[0]: + prefix_pieces = await reader.collect_window( + plan.local_window[0], first_blocked_pid + ) + for pid in range(plan.local_window[0], first_blocked_pid): await _emit_partition( context, ref_ir, @@ -637,37 +616,35 @@ async def finish_sends() -> None: ) # Read only far enough to satisfy destination liabilities before receiving. - if owed_remote_pids: - pieces = await reader.collect_window(*owed_remote_range) + if plan.owed_remote_pids: + pieces = await reader.collect_window(*plan.owed_remote_range) for pid, store in pieces.items(): - owner = _contiguous_owner(pid, comm.nranks, npartitions) + owner = _contiguous_owner(pid, comm.nranks, plan.npartitions) if owner == comm.rank: local_pieces[pid] = store continue - for msg in store: - buffered_remote_pids.add(pid) - await _send_remote_piece( - context, - comm, - exchange, - npartitions, - pid, - TableChunk.from_message(msg, br=context.br()), - ) - for pid in owed_remote_pids: - if pid not in buffered_remote_pids: - await _send_remote_marker( + buffered_remote_pids.add(pid) + await _send_remote_store( context, comm, exchange, - npartitions, + plan.npartitions, pid, + store, ) + await _send_missing_remote_markers( + context, + comm, + exchange, + plan.npartitions, + plan.owed_remote_pids, + buffered_remote_pids, + ) await exchange.insert_finished(context) pieces_by_source: dict[int, dict[int, ChunkStore]] = {} pieces_by_source[comm.rank] = local_pieces - for source_rank in remote_sources: + for source_rank in plan.remote_sources: remote_pieces: dict[int, ChunkStore] = {} stream = context.br().stream_pool.get_stream() for packed in exchange.extract(source_rank): @@ -678,8 +655,8 @@ async def finish_sends() -> None: _store_chunk(context, remote_pieces, pid, chunk) pieces_by_source[source_rank] = remote_pieces - for pid in range(first_blocked_pid, local_window[1]): - pid_sources = _sources_for_pid(source_ranges, pid) + for pid in range(first_blocked_pid, plan.local_window[1]): + pid_sources = _sources_for_pid(plan.source_ranges, pid) if comm.rank in pid_sources and pid not in local_pieces: local_pieces.update( { @@ -687,7 +664,7 @@ async def finish_sends() -> None: for piece_pid, store in ( await reader.collect_window(pid, pid + 1) ).items() - if _contiguous_owner(piece_pid, comm.nranks, npartitions) + if _contiguous_owner(piece_pid, comm.nranks, plan.npartitions) == comm.rank } ) From 07346d71e37b62230a375f29c59fde658a969dff Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 13 Jul 2026 07:21:45 -0700 Subject: [PATCH 19/26] simplify/combine codepaths --- .../actor_graph/collectives/ordering.py | 167 ++++++------------ 1 file changed, 56 insertions(+), 111 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index f1349ec0d061..48ec3d471c98 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -519,77 +519,18 @@ async def _adjust_ordering_impl( plan = _make_routing_plan(context, comm, input_ordering, output_ordering) reader = _OutputPieceReader(context, ch_in, plan.boundary_chunk, output_ordering) - # Ranks with no incoming dependency can stream local output immediately - # while sending remote-owned pieces as they are encountered. - if not plan.remote_sources: - if plan.remote_destinations: - assert collective_id is not None - exchange = SparseAlltoall( - context, - comm, - collective_id, - srcs=[], - dsts=plan.remote_destinations, - ) - else: - exchange = None - - streamed_remote_pids: set[int] = set() - start, stop = plan.local_window - if plan.owed_remote_pids: - start = min(start, plan.owed_remote_range[0]) - stop = max(stop, plan.owed_remote_range[1]) - for pid in range(start, stop): - pieces = await reader.collect_window(pid, pid + 1) - owner = _contiguous_owner(pid, comm.nranks, plan.npartitions) - if ( - owner == comm.rank - and plan.local_window[0] <= pid < plan.local_window[1] - ): - await _emit_partition( - context, - ref_ir, - ir_context, - ch_out, - pid, - pieces.get(pid), - ) - continue - if exchange is not None and pid in pieces: - streamed_remote_pids.add(pid) - await _send_remote_store( - context, - comm, - exchange, - plan.npartitions, - pid, - pieces[pid], - ) - if exchange is not None: - await _send_missing_remote_markers( - context, - comm, - exchange, - plan.npartitions, - plan.owed_remote_pids, - streamed_remote_pids, - ) - await exchange.insert_finished(context) - await reader.drain_remaining() - await ch_out.drain(context) - return - - assert collective_id is not None - exchange = SparseAlltoall( - context, - comm, - collective_id, - srcs=plan.remote_sources, - dsts=plan.remote_destinations, - ) - + exchange = None + if plan.remote_sources or plan.remote_destinations: + assert collective_id is not None + exchange = SparseAlltoall( + context, + comm, + collective_id, + srcs=plan.remote_sources, + dsts=plan.remote_destinations, + ) local_pieces: dict[int, ChunkStore] = {} - buffered_remote_pids: set[int] = set() + sent_remote_pids: set[int] = set() first_blocked_pid = next( ( pid @@ -601,59 +542,63 @@ async def _adjust_ordering_impl( ), plan.local_window[1], ) - if first_blocked_pid > plan.local_window[0]: - prefix_pieces = await reader.collect_window( - plan.local_window[0], first_blocked_pid - ) - for pid in range(plan.local_window[0], first_blocked_pid): - await _emit_partition( - context, - ref_ir, - ir_context, - ch_out, - pid, - prefix_pieces.get(pid), - ) - # Read only far enough to satisfy destination liabilities before receiving. + pre_start, pre_stop = plan.local_window[0], first_blocked_pid if plan.owed_remote_pids: - pieces = await reader.collect_window(*plan.owed_remote_range) - for pid, store in pieces.items(): - owner = _contiguous_owner(pid, comm.nranks, plan.npartitions) - if owner == comm.rank: - local_pieces[pid] = store - continue - buffered_remote_pids.add(pid) + pre_start = min(pre_start, plan.owed_remote_range[0]) + pre_stop = max(pre_stop, plan.owed_remote_range[1]) + + # Before receiving remote pieces, emit the local-only prefix and send all + # remote-owned pieces this rank is responsible for. + for pid in range(pre_start, pre_stop): + pieces = await reader.collect_window(pid, pid + 1) + owner = _contiguous_owner(pid, comm.nranks, plan.npartitions) + if exchange is not None and owner != comm.rank and pid in pieces: + sent_remote_pids.add(pid) await _send_remote_store( context, comm, exchange, plan.npartitions, pid, - store, + pieces[pid], ) - await _send_missing_remote_markers( - context, - comm, - exchange, - plan.npartitions, - plan.owed_remote_pids, - buffered_remote_pids, - ) - await exchange.insert_finished(context) + elif owner == comm.rank and pid in pieces: + local_pieces[pid] = pieces[pid] + if plan.local_window[0] <= pid < first_blocked_pid: + await _emit_partition( + context, + ref_ir, + ir_context, + ch_out, + pid, + local_pieces.pop(pid, None), + ) + + if exchange is not None: + await _send_missing_remote_markers( + context, + comm, + exchange, + plan.npartitions, + plan.owed_remote_pids, + sent_remote_pids, + ) + await exchange.insert_finished(context) pieces_by_source: dict[int, dict[int, ChunkStore]] = {} pieces_by_source[comm.rank] = local_pieces - for source_rank in plan.remote_sources: - remote_pieces: dict[int, ChunkStore] = {} - stream = context.br().stream_pool.get_stream() - for packed in exchange.extract(source_rank): - remote_piece = _unpack_remote_piece(packed, stream, context.br()) - if remote_piece is None: - continue - pid, chunk = remote_piece - _store_chunk(context, remote_pieces, pid, chunk) - pieces_by_source[source_rank] = remote_pieces + if exchange is not None: + for source_rank in plan.remote_sources: + remote_pieces: dict[int, ChunkStore] = {} + stream = context.br().stream_pool.get_stream() + for packed in exchange.extract(source_rank): + remote_piece = _unpack_remote_piece(packed, stream, context.br()) + if remote_piece is None: + continue + pid, chunk = remote_piece + _store_chunk(context, remote_pieces, pid, chunk) + pieces_by_source[source_rank] = remote_pieces for pid in range(first_blocked_pid, plan.local_window[1]): pid_sources = _sources_for_pid(plan.source_ranges, pid) From e36bcb43eba5ee54d5dc6466ba41361908bfd99d Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 13 Jul 2026 07:46:05 -0700 Subject: [PATCH 20/26] update test coverage --- .../actor_graph/collectives/ordering.py | 63 ++++++++++------ .../tests/streaming/test_adjust_ordering.py | 74 ++++++++++++++++++- 2 files changed, 111 insertions(+), 26 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 48ec3d471c98..8d03c9724b19 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -42,6 +42,7 @@ _PID_DTYPE = DataType(pl.Int32()) _PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32) +_MARKER_PLC_DTYPE = plc.DataType(plc.TypeId.BOOL8) @dataclass(frozen=True) @@ -110,14 +111,29 @@ def _split_points( ) -def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Table: - """Append a temporary target-partition-id column to *table*.""" +def _append_remote_control_columns( + table: plc.Table, + pid: int, + stream: Stream, + *, + is_marker: bool, + nrows: int | None = None, +) -> plc.Table: + """Append temporary pid and marker columns to one remote message.""" + nrows = table.num_rows() if nrows is None else nrows + # The marker column keeps empty-control messages distinct from future + # zero-column payloads. pid_col = plc.Column.from_scalar( plc.Scalar.from_py(pid, _PID_PLC_DTYPE, stream=stream), - table.num_rows(), + nrows, stream=stream, ) - return plc.Table([*table.columns(), pid_col]) + marker_col = plc.Column.from_scalar( + plc.Scalar.from_py(is_marker, _MARKER_PLC_DTYPE, stream=stream), + nrows, + stream=stream, + ) + return plc.Table([*table.columns(), pid_col, marker_col]) def _boundary_search_positions( @@ -213,18 +229,16 @@ def _unpack_remote_piece( table = unpack_and_concat([packed], stream=stream, br=br) if table.num_rows() == 0: return None - *payload_cols, pid_col = table.columns() - pid = int( - DataFrame.from_table( - plc.Table([pid_col]), - ["pid"], - [_PID_DTYPE], - stream, - ) - .to_polars() - .item(0, 0) - ) - if not payload_cols: + *payload_cols, pid_col, marker_col = table.columns() + control = DataFrame.from_table( + plc.Table([pid_col, marker_col]), + ["pid", "is_marker"], + [_PID_DTYPE, DataType(pl.Boolean())], + stream, + ).to_polars() + pid = int(control["pid"].item(0)) + is_marker = bool(control["is_marker"].item(0)) + if is_marker: return None payload = plc.Table(payload_cols).copy(stream=stream, mr=br.device_mr) return pid, TableChunk.from_pylibcudf_table( @@ -366,7 +380,9 @@ async def _send_remote_piece( _contiguous_owner(pid, comm.nranks, npartitions), packed_data_from_cudf_packed_columns( pack( - _append_partition_id(chunk.table_view(), pid, stream), + _append_remote_control_columns( + chunk.table_view(), pid, stream, is_marker=False + ), stream, mr=context.br().device_mr, ), @@ -388,15 +404,16 @@ async def _send_remote_marker( context.br().stream_pool.get_stream, upstreams=(), ) as stream: - pid_col = plc.Column.from_scalar( - plc.Scalar.from_py(pid, _PID_PLC_DTYPE, stream=stream), - 1, - stream=stream, - ) exchange.insert( _contiguous_owner(pid, comm.nranks, npartitions), packed_data_from_cudf_packed_columns( - pack(plc.Table([pid_col]), stream, mr=context.br().device_mr), + pack( + _append_remote_control_columns( + plc.Table([]), pid, stream, is_marker=True, nrows=1 + ), + stream, + mr=context.br().device_mr, + ), stream, context.br(), ), diff --git a/python/cudf_polars/tests/streaming/test_adjust_ordering.py b/python/cudf_polars/tests/streaming/test_adjust_ordering.py index 86158044163b..4d5561290f4a 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_ordering.py +++ b/python/cudf_polars/tests/streaming/test_adjust_ordering.py @@ -51,6 +51,8 @@ def _make_ordering( boundary: _Boundary | list[_Boundary], *, key_indices: tuple[int, ...] = (0,), + order: plc.types.Order = plc.types.Order.ASCENDING, + null_order: plc.types.NullOrder = plc.types.NullOrder.BEFORE, strict: bool = True, stream: Stream, ) -> Ordering: @@ -73,8 +75,8 @@ def _make_ordering( [ OrderKey( index, - plc.types.Order.ASCENDING, - plc.types.NullOrder.BEFORE, + order, + null_order, ) for index in key_indices ], @@ -267,10 +269,12 @@ def test_adjust_ordering_requires_collective_id( (5, {0: {0: [0, 1, 2, 3, 4]}, 1: {1: [5, 6, 7]}}), ], ) +@pytest.mark.parametrize("input_strict", [True, False]) def test_adjust_ordering_sparse_boundary_shift( spmd_engine: SPMDEngine, target_boundary: int, expected: _ExpectedByRank, + input_strict: bool, # noqa: FBT001 ) -> None: context = spmd_engine.context comm = spmd_engine.comm @@ -280,7 +284,13 @@ def test_adjust_ordering_sparse_boundary_shift( keys = list(range(4)) if comm.rank == 0 else list(range(4, 8)) stream = context.br().stream_pool.get_stream() # Input sorted on (key, val) is also sorted on the target key prefix. - input_ordering = _make_ordering(context, (4, 4), key_indices=(0, 1), stream=stream) + input_ordering = _make_ordering( + context, + (4, 4), + key_indices=(0, 1), + strict=input_strict, + stream=stream, + ) output_ordering = _make_ordering(context, target_boundary, stream=stream) with reserve_op_id() as op_id: @@ -496,3 +506,61 @@ def test_adjust_ordering_multi_chunk_input(spmd_engine: SPMDEngine) -> None: ) _assert_partition_output(output, {0: [0, 1, 2, 3], 1: [4, 5, 6, 7]}) + + +@pytest.mark.spmd +@pytest.mark.parametrize( + "order,null_order,keys,expected", + [ + ( + plc.types.Order.DESCENDING, + plc.types.NullOrder.BEFORE, + list(range(7, -1, -1)), + {0: [7, 6, 5], 1: [4, 3, 2, 1, 0]}, + ), + ( + plc.types.Order.ASCENDING, + plc.types.NullOrder.AFTER, + list(range(8)), + {0: [0, 1, 2, 3], 1: [4, 5, 6, 7]}, + ), + ], +) +def test_adjust_ordering_respects_order_key_metadata( + spmd_engine: SPMDEngine, + order: plc.types.Order, + null_order: plc.types.NullOrder, + keys: list[int], + expected: _ExpectedPartitions, +) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks != 1: + pytest.skip("This test covers local order-key metadata variants.") + + stream = context.br().stream_pool.get_stream() + input_ordering = _make_ordering( + context, + 4, + order=order, + null_order=null_order, + stream=stream, + ) + output_ordering = _make_ordering( + context, + 4, + order=order, + null_order=null_order, + stream=stream, + ) + output_by_pid = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(keys), + input_ordering, + output_ordering, + ) + ) + + _assert_partition_output(output_by_pid, expected) From bf0d354a753c7e001b83009dda106b9bbfe98191 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 13 Jul 2026 08:17:08 -0700 Subject: [PATCH 21/26] remove empty-marker machinery --- .../actor_graph/collectives/ordering.py | 111 +++--------------- 1 file changed, 18 insertions(+), 93 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 8d03c9724b19..e11f20aed2b6 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -42,7 +42,6 @@ _PID_DTYPE = DataType(pl.Int32()) _PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32) -_MARKER_PLC_DTYPE = plc.DataType(plc.TypeId.BOOL8) @dataclass(frozen=True) @@ -111,29 +110,14 @@ def _split_points( ) -def _append_remote_control_columns( - table: plc.Table, - pid: int, - stream: Stream, - *, - is_marker: bool, - nrows: int | None = None, -) -> plc.Table: - """Append temporary pid and marker columns to one remote message.""" - nrows = table.num_rows() if nrows is None else nrows - # The marker column keeps empty-control messages distinct from future - # zero-column payloads. +def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Table: + """Append a temporary target-partition-id column to one remote piece.""" pid_col = plc.Column.from_scalar( plc.Scalar.from_py(pid, _PID_PLC_DTYPE, stream=stream), - nrows, - stream=stream, - ) - marker_col = plc.Column.from_scalar( - plc.Scalar.from_py(is_marker, _MARKER_PLC_DTYPE, stream=stream), - nrows, + table.num_rows(), stream=stream, ) - return plc.Table([*table.columns(), pid_col, marker_col]) + return plc.Table([*table.columns(), pid_col]) def _boundary_search_positions( @@ -224,22 +208,20 @@ def _unpack_remote_piece( packed: PackedData, stream: Stream, br: BufferResource, -) -> tuple[int, TableChunk] | None: +) -> tuple[int, TableChunk]: """Unpack one remote piece and recover its temporary target partition ID.""" table = unpack_and_concat([packed], stream=stream, br=br) - if table.num_rows() == 0: - return None - *payload_cols, pid_col, marker_col = table.columns() - control = DataFrame.from_table( - plc.Table([pid_col, marker_col]), - ["pid", "is_marker"], - [_PID_DTYPE, DataType(pl.Boolean())], - stream, - ).to_polars() - pid = int(control["pid"].item(0)) - is_marker = bool(control["is_marker"].item(0)) - if is_marker: - return None + *payload_cols, pid_col = table.columns() + pid = int( + DataFrame.from_table( + plc.Table([pid_col]), + ["pid"], + [_PID_DTYPE], + stream, + ) + .to_polars() + .item(0, 0) + ) payload = plc.Table(payload_cols).copy(stream=stream, mr=br.device_mr) return pid, TableChunk.from_pylibcudf_table( payload, @@ -380,37 +362,7 @@ async def _send_remote_piece( _contiguous_owner(pid, comm.nranks, npartitions), packed_data_from_cudf_packed_columns( pack( - _append_remote_control_columns( - chunk.table_view(), pid, stream, is_marker=False - ), - stream, - mr=context.br().device_mr, - ), - stream, - context.br(), - ), - ) - - -async def _send_remote_marker( - context: Context, - comm: Communicator, - exchange: SparseAlltoall, - npartitions: int, - pid: int, -) -> None: - """Send a no-payload marker for an empty remote-owned partition.""" - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(), - ) as stream: - exchange.insert( - _contiguous_owner(pid, comm.nranks, npartitions), - packed_data_from_cudf_packed_columns( - pack( - _append_remote_control_columns( - plc.Table([]), pid, stream, is_marker=True, nrows=1 - ), + _append_partition_id(chunk.table_view(), pid, stream), stream, mr=context.br().device_mr, ), @@ -420,20 +372,6 @@ async def _send_remote_marker( ) -async def _send_missing_remote_markers( - context: Context, - comm: Communicator, - exchange: SparseAlltoall, - npartitions: int, - owed_remote_pids: list[int], - sent_remote_pids: set[int], -) -> None: - """Send empty markers for remote-owned partitions with no payload.""" - for pid in owed_remote_pids: - if pid not in sent_remote_pids: - await _send_remote_marker(context, comm, exchange, npartitions, pid) - - async def _emit_partition( context: Context, ref_ir: IR, @@ -547,7 +485,6 @@ async def _adjust_ordering_impl( dsts=plan.remote_destinations, ) local_pieces: dict[int, ChunkStore] = {} - sent_remote_pids: set[int] = set() first_blocked_pid = next( ( pid @@ -571,7 +508,6 @@ async def _adjust_ordering_impl( pieces = await reader.collect_window(pid, pid + 1) owner = _contiguous_owner(pid, comm.nranks, plan.npartitions) if exchange is not None and owner != comm.rank and pid in pieces: - sent_remote_pids.add(pid) await _send_remote_store( context, comm, @@ -593,14 +529,6 @@ async def _adjust_ordering_impl( ) if exchange is not None: - await _send_missing_remote_markers( - context, - comm, - exchange, - plan.npartitions, - plan.owed_remote_pids, - sent_remote_pids, - ) await exchange.insert_finished(context) pieces_by_source: dict[int, dict[int, ChunkStore]] = {} @@ -610,10 +538,7 @@ async def _adjust_ordering_impl( remote_pieces: dict[int, ChunkStore] = {} stream = context.br().stream_pool.get_stream() for packed in exchange.extract(source_rank): - remote_piece = _unpack_remote_piece(packed, stream, context.br()) - if remote_piece is None: - continue - pid, chunk = remote_piece + pid, chunk = _unpack_remote_piece(packed, stream, context.br()) _store_chunk(context, remote_pieces, pid, chunk) pieces_by_source[source_rank] = remote_pieces From bfe2d04d4acc7e5392c232e466296ea1b82186f2 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Mon, 13 Jul 2026 11:44:03 -0700 Subject: [PATCH 22/26] end-to-end cleanup --- .../actor_graph/collectives/ordering.py | 115 ++++++++++-------- .../tests/streaming/test_adjust_ordering.py | 45 ++++++- 2 files changed, 110 insertions(+), 50 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index e11f20aed2b6..026a0a7f4d3d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -46,14 +46,22 @@ @dataclass(frozen=True) class _RoutingPlan: + """Static ownership and exchange plan for one ordering adjustment.""" + npartitions: int - boundary_chunk: TableChunk + """Number of output partitions implied by the target ordering.""" local_window: tuple[int, int] + """Half-open output-partition range owned by this rank.""" source_ranges: list[tuple[int, int]] + """Half-open output-partition range each source rank may contribute to.""" remote_sources: list[int] + """Remote ranks that may contribute to this rank's local output window.""" owed_remote_pids: list[int] + """Remote-owned output partitions this rank may contribute to.""" remote_destinations: list[int] + """Remote ranks that may receive this rank's data.""" owed_remote_range: tuple[int, int] + """Half-open range spanning owed remote pids, or empty when none are owed.""" def _contiguous_owner(pid: int, nranks: int, npartitions: int) -> int: @@ -246,8 +254,8 @@ def _copy_to_owned_chunk( ) -class _OutputPieceReader: - """Read input only far enough to materialize requested output windows.""" +class _OutputPartitionBuffer: + """Incrementally split and buffer pieces by target output partition.""" def __init__( self, @@ -263,11 +271,21 @@ def __init__( self.pending: dict[int, ChunkStore] = {} self.input_done = False - async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: - """Return all locally-read pieces for output pids in ``[start, stop)``.""" - if start >= stop: - return {} - while not self.input_done and not any(pid >= stop for pid in self.pending): + async def collect_output_partition(self, pid: int) -> ChunkStore | None: + """ + Return buffered pieces for one output pid. + + Notes + ----- + This reads ordered input chunks until the requested output partition + is complete. All pieces produced by those chunks are held in a + spillable container, and pieces for later output partitions remain + cached for later calls. + """ + stop = pid + 1 + while not self.input_done and not any( + pending_pid >= stop for pending_pid in self.pending + ): msg = await self.ch_in.recv(self.context) if msg is None: self.input_done = True @@ -288,7 +306,7 @@ async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: self.output_ordering, stream, ) - for pid, piece in enumerate( + for piece_pid, piece in enumerate( plc.copying.split(table, splits, stream=stream) ): if piece.num_rows() == 0: @@ -296,22 +314,24 @@ async def collect_window(self, start: int, stop: int) -> dict[int, ChunkStore]: _store_chunk( self.context, self.pending, - pid, + piece_pid, _copy_to_owned_chunk(piece, stream, self.context.br()), ) + return self.pending.pop(pid, None) - out = { - pid: self.pending.pop(pid) - for pid in list(self.pending) - if start <= pid < stop - } - return dict(sorted(out.items())) - - async def drain_remaining(self) -> None: - """Consume input that cannot affect any remaining output partition.""" - while not self.input_done: - if await self.ch_in.recv(self.context) is None: - self.input_done = True + async def assert_input_drained(self) -> None: + """Verify that no buffered or unread input remains.""" + if self.pending: + raise RuntimeError( + "adjust_ordering left buffered data after all output was emitted." + ) + if not self.input_done: + while (msg := await self.ch_in.recv(self.context)) is not None: + if TableChunk.from_message(msg, br=self.context.br()).shape[0] > 0: + raise RuntimeError( + "adjust_ordering left unread input after all output was emitted" + ) + self.input_done = True def _store_chunk( @@ -399,10 +419,10 @@ def _make_routing_plan( comm: Communicator, input_ordering: Ordering, output_ordering: Ordering, + output_boundaries: TableChunk, ) -> _RoutingPlan: """Compute local ownership and sparse-exchange obligations.""" npartitions = output_ordering.num_boundaries + 1 - boundary_chunk = output_ordering.get_boundaries(context.br()) local_window = _partition_range(comm.rank, comm.nranks, npartitions) if comm.nranks == 1: @@ -411,11 +431,11 @@ def _make_routing_plan( input_boundary_chunk = input_ordering.get_boundaries(context.br()) with stream_ordered_after( context.br().stream_pool.get_stream, - upstreams=(input_boundary_chunk.stream, boundary_chunk.stream), + upstreams=(input_boundary_chunk.stream, output_boundaries.stream), ) as stream: lower_positions, upper_positions = _boundary_search_positions( input_boundary_chunk.table_view(), - boundary_chunk.table_view(), + output_boundaries.table_view(), output_ordering, stream, ) @@ -449,7 +469,6 @@ def _make_routing_plan( ) return _RoutingPlan( npartitions=npartitions, - boundary_chunk=boundary_chunk, local_window=local_window, source_ranges=source_ranges, remote_sources=remote_sources, @@ -471,8 +490,11 @@ async def _adjust_ordering_impl( collective_id: int | None, ) -> None: """Adjust ordering while using exchange only for remote dependencies.""" - plan = _make_routing_plan(context, comm, input_ordering, output_ordering) - reader = _OutputPieceReader(context, ch_in, plan.boundary_chunk, output_ordering) + output_boundaries = output_ordering.get_boundaries(context.br()) + plan = _make_routing_plan( + context, comm, input_ordering, output_ordering, output_boundaries + ) + buffer = _OutputPartitionBuffer(context, ch_in, output_boundaries, output_ordering) exchange = None if plan.remote_sources or plan.remote_destinations: @@ -484,7 +506,7 @@ async def _adjust_ordering_impl( srcs=plan.remote_sources, dsts=plan.remote_destinations, ) - local_pieces: dict[int, ChunkStore] = {} + first_blocked_pid = next( ( pid @@ -497,27 +519,29 @@ async def _adjust_ordering_impl( plan.local_window[1], ) + # Before receiving remote pieces, emit the local-only prefix and send all + # remote-owned pieces this rank is responsible for. pre_start, pre_stop = plan.local_window[0], first_blocked_pid if plan.owed_remote_pids: pre_start = min(pre_start, plan.owed_remote_range[0]) pre_stop = max(pre_stop, plan.owed_remote_range[1]) - # Before receiving remote pieces, emit the local-only prefix and send all - # remote-owned pieces this rank is responsible for. + pieces_by_source: dict[int, dict[int, ChunkStore]] = {comm.rank: {}} + local_pieces = pieces_by_source[comm.rank] for pid in range(pre_start, pre_stop): - pieces = await reader.collect_window(pid, pid + 1) + piece = await buffer.collect_output_partition(pid) owner = _contiguous_owner(pid, comm.nranks, plan.npartitions) - if exchange is not None and owner != comm.rank and pid in pieces: + if exchange is not None and owner != comm.rank and piece is not None: await _send_remote_store( context, comm, exchange, plan.npartitions, pid, - pieces[pid], + piece, ) - elif owner == comm.rank and pid in pieces: - local_pieces[pid] = pieces[pid] + elif owner == comm.rank and piece is not None: + local_pieces[pid] = piece if plan.local_window[0] <= pid < first_blocked_pid: await _emit_partition( context, @@ -531,8 +555,7 @@ async def _adjust_ordering_impl( if exchange is not None: await exchange.insert_finished(context) - pieces_by_source: dict[int, dict[int, ChunkStore]] = {} - pieces_by_source[comm.rank] = local_pieces + # Collect remote pieces from other ranks. if exchange is not None: for source_rank in plan.remote_sources: remote_pieces: dict[int, ChunkStore] = {} @@ -542,19 +565,13 @@ async def _adjust_ordering_impl( _store_chunk(context, remote_pieces, pid, chunk) pieces_by_source[source_rank] = remote_pieces + # Collect and emit the remaining output partitions for this rank. for pid in range(first_blocked_pid, plan.local_window[1]): pid_sources = _sources_for_pid(plan.source_ranges, pid) if comm.rank in pid_sources and pid not in local_pieces: - local_pieces.update( - { - piece_pid: store - for piece_pid, store in ( - await reader.collect_window(pid, pid + 1) - ).items() - if _contiguous_owner(piece_pid, comm.nranks, plan.npartitions) - == comm.rank - } - ) + piece = await buffer.collect_output_partition(pid) + if piece is not None: + local_pieces[pid] = piece chunks: list[TableChunk] = [] for source_rank in pid_sources: stores = pieces_by_source.get(source_rank) @@ -572,7 +589,7 @@ async def _adjust_ordering_impl( else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) ) await ch_out.send(context, Message(pid, chunk)) - await reader.drain_remaining() + await buffer.assert_input_drained() await ch_out.drain(context) diff --git a/python/cudf_polars/tests/streaming/test_adjust_ordering.py b/python/cudf_polars/tests/streaming/test_adjust_ordering.py index 4d5561290f4a..c35cd608d24c 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_ordering.py +++ b/python/cudf_polars/tests/streaming/test_adjust_ordering.py @@ -308,6 +308,49 @@ def test_adjust_ordering_sparse_boundary_shift( _assert_partition_output(output, expected[comm.rank]) +@pytest.mark.spmd +def test_adjust_ordering_descending_sparse_boundary_shift( + spmd_engine: SPMDEngine, +) -> None: + context = spmd_engine.context + comm = spmd_engine.comm + if comm.nranks != 2: + pytest.skip("This test expects exactly two ranks.") + + keys = list(range(7, 3, -1)) if comm.rank == 0 else list(range(3, -1, -1)) + stream = context.br().stream_pool.get_stream() + input_ordering = _make_ordering( + context, + 3, + order=plc.types.Order.DESCENDING, + stream=stream, + ) + output_ordering = _make_ordering( + context, + 5, + order=plc.types.Order.DESCENDING, + stream=stream, + ) + + with reserve_op_id() as op_id: + output = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(keys), + input_ordering, + output_ordering, + collective_id=op_id, + ) + ) + + expected = { + 0: {0: [7, 6]}, + 1: {1: [5, 4, 3, 2, 1, 0]}, + }[comm.rank] + _assert_partition_output(output, expected) + + @pytest.mark.spmd def test_adjust_ordering_emits_empty_owned_partitions( spmd_engine: SPMDEngine, @@ -499,7 +542,7 @@ def test_adjust_ordering_multi_chunk_input(spmd_engine: SPMDEngine) -> None: _adjust_and_collect( context, comm, - [_frame([0, 1]), _frame([2, 3, 4, 5]), _frame([6, 7])], + [_frame([0, 1]), _frame([2, 3, 4, 5]), _frame([6, 7]), _frame([])], input_ordering, output_ordering, ) From 536348c91781725905c8388448dcde5d2085f374 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Tue, 14 Jul 2026 14:42:21 -0700 Subject: [PATCH 23/26] partial review pass --- .../actor_graph/collectives/ordering.py | 245 ++++++++++-------- 1 file changed, 137 insertions(+), 108 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 026a0a7f4d3d..9fde72a867b0 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -42,6 +42,7 @@ _PID_DTYPE = DataType(pl.Int32()) _PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32) +_PartitionRange = tuple[int, int] @dataclass(frozen=True) @@ -50,9 +51,9 @@ class _RoutingPlan: npartitions: int """Number of output partitions implied by the target ordering.""" - local_window: tuple[int, int] + local_window: _PartitionRange """Half-open output-partition range owned by this rank.""" - source_ranges: list[tuple[int, int]] + source_ranges: list[_PartitionRange] """Half-open output-partition range each source rank may contribute to.""" remote_sources: list[int] """Remote ranks that may contribute to this rank's local output window.""" @@ -60,16 +61,86 @@ class _RoutingPlan: """Remote-owned output partitions this rank may contribute to.""" remote_destinations: list[int] """Remote ranks that may receive this rank's data.""" - owed_remote_range: tuple[int, int] + owed_remote_range: _PartitionRange """Half-open range spanning owed remote pids, or empty when none are owed.""" + @classmethod + def from_orderings( + cls, + context: Context, + comm: Communicator, + input_ordering: Ordering, + output_ordering: Ordering, + output_boundaries: TableChunk, + ) -> _RoutingPlan: + """Compute local ownership and sparse-exchange obligations.""" + npartitions = output_ordering.num_boundaries + 1 + local_window = _partition_range(comm.rank, comm.nranks, npartitions) + + if comm.nranks == 1: + source_ranges = [(0, npartitions)] + else: + input_boundary_chunk = input_ordering.get_boundaries(context.br()) + with stream_ordered_after( + context.br().stream_pool.get_stream, + upstreams=(input_boundary_chunk.stream, output_boundaries.stream), + ) as stream: + lower_positions, upper_positions = _boundary_search_positions( + input_boundary_chunk.table_view(), + output_boundaries.table_view(), + output_ordering, + stream, + ) + source_ranges = [ + _source_output_range( + source_rank, + comm.nranks, + input_ordering, + output_ordering, + lower_positions, + upper_positions, + ) + for source_rank in range(comm.nranks) + ] + local_source_range = source_ranges[comm.rank] + remote_sources = [ + source_rank + for source_rank, source_range in enumerate(source_ranges) + if source_rank != comm.rank and _ranges_overlap(source_range, local_window) + ] + owed_remote_pids = [ + pid + for pid in range(*local_source_range) + if _contiguous_owner(pid, comm.nranks, npartitions) != comm.rank + ] + remote_destinations = sorted( + { + _contiguous_owner(pid, comm.nranks, npartitions) + for pid in owed_remote_pids + } + ) + owed_remote_range = ( + (owed_remote_pids[0], owed_remote_pids[-1] + 1) + if owed_remote_pids + else (0, 0) + ) + return cls( + npartitions=npartitions, + local_window=local_window, + source_ranges=source_ranges, + remote_sources=remote_sources, + owed_remote_pids=owed_remote_pids, + remote_destinations=remote_destinations, + owed_remote_range=owed_remote_range, + ) + def _contiguous_owner(pid: int, nranks: int, npartitions: int) -> int: """Return the rank owning *pid* under contiguous partition assignment.""" return pid * nranks // npartitions -def _partition_range(rank: int, nranks: int, npartitions: int) -> tuple[int, int]: +def _partition_range(rank: int, nranks: int, npartitions: int) -> _PartitionRange: """Return the half-open partition ID range owned by *rank*.""" return ( (rank * npartitions + nranks - 1) // nranks, @@ -95,7 +166,12 @@ def _split_points( ordering: Ordering, stream: Stream, ) -> list[int]: - """Return row split points that partition *table* by *ordering* boundaries.""" + """ + Return row split points that partition *table* by *ordering* boundaries. + + ``table`` and ``boundary_table`` must be valid on ``stream``. The returned + host list is materialized from device data on ``stream``. + """ if boundary_table.num_rows() == 0: return [] key_table = plc.Table([table.columns()[key.column_index] for key in ordering.keys]) @@ -119,7 +195,12 @@ def _split_points( def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Table: - """Append a temporary target-partition-id column to one remote piece.""" + """ + Append the target partition ID before packing a remote table piece. + + SparseAlltoall exchanges packed table payloads, so the partition ID travels + as a temporary column and is stripped by the receiver. + """ pid_col = plc.Column.from_scalar( plc.Scalar.from_py(pid, _PID_PLC_DTYPE, stream=stream), table.num_rows(), @@ -134,7 +215,12 @@ def _boundary_search_positions( output_ordering: Ordering, stream: Stream, ) -> tuple[list[int], list[int]]: - """Search output boundary positions for projected input boundary rows.""" + """ + Search output boundary positions for projected input boundary rows. + + Boundary tables must be valid on ``stream``. The returned host lists are + materialized from device data on ``stream``. + """ if input_boundary_table.num_rows() == 0: return [], [] prefix_len = len(output_ordering.keys) @@ -172,7 +258,7 @@ def _source_output_range( output_ordering: Ordering, lower_positions: list[int], upper_positions: list[int], -) -> tuple[int, int]: +) -> _PartitionRange: """Return the half-open output partition range touched by a source rank.""" input_npartitions = input_ordering.num_boundaries + 1 output_npartitions = output_ordering.num_boundaries + 1 @@ -195,13 +281,13 @@ def _source_output_range( return output_start, output_stop -def _ranges_overlap(left: tuple[int, int], right: tuple[int, int]) -> bool: +def _ranges_overlap(left: _PartitionRange, right: _PartitionRange) -> bool: """Return whether two half-open integer ranges overlap.""" return max(left[0], right[0]) < min(left[1], right[1]) def _sources_for_pid( - source_ranges: list[tuple[int, int]], + source_ranges: list[_PartitionRange], pid: int, ) -> list[int]: """Return source ranks that may contribute to one output partition.""" @@ -325,12 +411,19 @@ async def assert_input_drained(self) -> None: raise RuntimeError( "adjust_ordering left buffered data after all output was emitted." ) - if not self.input_done: - while (msg := await self.ch_in.recv(self.context)) is not None: - if TableChunk.from_message(msg, br=self.context.br()).shape[0] > 0: - raise RuntimeError( - "adjust_ordering left unread input after all output was emitted" - ) + if ( + not self.input_done + and (msg := await self.ch_in.recv(self.context)) is not None + ): + rows = ( + TableChunk.from_message(msg, br=self.context.br()) + .table_view() + .num_rows() + ) + raise RuntimeError( + "adjust_ordering left unread input after all output was emitted " + f"({rows} rows)" + ) self.input_done = True @@ -374,27 +467,24 @@ async def _send_remote_piece( chunk: TableChunk, ) -> None: """Send one output-partition piece to its remote owner.""" - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(chunk.stream,), - ) as stream: - exchange.insert( - _contiguous_owner(pid, comm.nranks, npartitions), - packed_data_from_cudf_packed_columns( - pack( - _append_partition_id(chunk.table_view(), pid, stream), - stream, - mr=context.br().device_mr, - ), + stream = chunk.stream + exchange.insert( + _contiguous_owner(pid, comm.nranks, npartitions), + packed_data_from_cudf_packed_columns( + pack( + _append_partition_id(chunk.table_view(), pid, stream), stream, - context.br(), + mr=context.br().device_mr, ), - ) + stream, + context.br(), + ), + ) async def _emit_partition( context: Context, - ref_ir: IR, + schema_ir: IR, ir_context: IRExecutionContext, ch_out: Channel[TableChunk], pid: int, @@ -407,81 +497,17 @@ async def _emit_partition( else [] ) chunk = ( - await concat_batch(chunks, context, ref_ir.schema, ir_context) + await concat_batch(chunks, context, schema_ir.schema, ir_context) if chunks - else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) + else empty_table_chunk(schema_ir, context, ir_context.get_cuda_stream()) ) await ch_out.send(context, Message(pid, chunk)) -def _make_routing_plan( - context: Context, - comm: Communicator, - input_ordering: Ordering, - output_ordering: Ordering, - output_boundaries: TableChunk, -) -> _RoutingPlan: - """Compute local ownership and sparse-exchange obligations.""" - npartitions = output_ordering.num_boundaries + 1 - local_window = _partition_range(comm.rank, comm.nranks, npartitions) - - if comm.nranks == 1: - source_ranges = [(0, npartitions)] - else: - input_boundary_chunk = input_ordering.get_boundaries(context.br()) - with stream_ordered_after( - context.br().stream_pool.get_stream, - upstreams=(input_boundary_chunk.stream, output_boundaries.stream), - ) as stream: - lower_positions, upper_positions = _boundary_search_positions( - input_boundary_chunk.table_view(), - output_boundaries.table_view(), - output_ordering, - stream, - ) - source_ranges = [ - _source_output_range( - source_rank, - comm.nranks, - input_ordering, - output_ordering, - lower_positions, - upper_positions, - ) - for source_rank in range(comm.nranks) - ] - local_source_range = source_ranges[comm.rank] - remote_sources = [ - source_rank - for source_rank, source_range in enumerate(source_ranges) - if source_rank != comm.rank and _ranges_overlap(source_range, local_window) - ] - owed_remote_pids = [ - pid - for pid in range(*local_source_range) - if _contiguous_owner(pid, comm.nranks, npartitions) != comm.rank - ] - remote_destinations = sorted( - {_contiguous_owner(pid, comm.nranks, npartitions) for pid in owed_remote_pids} - ) - owed_remote_range = ( - (owed_remote_pids[0], owed_remote_pids[-1] + 1) if owed_remote_pids else (0, 0) - ) - return _RoutingPlan( - npartitions=npartitions, - local_window=local_window, - source_ranges=source_ranges, - remote_sources=remote_sources, - owed_remote_pids=owed_remote_pids, - remote_destinations=remote_destinations, - owed_remote_range=owed_remote_range, - ) - - async def _adjust_ordering_impl( context: Context, comm: Communicator, - ref_ir: IR, + schema_ir: IR, ir_context: IRExecutionContext, ch_out: Channel[TableChunk], ch_in: Channel[TableChunk], @@ -491,7 +517,7 @@ async def _adjust_ordering_impl( ) -> None: """Adjust ordering while using exchange only for remote dependencies.""" output_boundaries = output_ordering.get_boundaries(context.br()) - plan = _make_routing_plan( + plan = _RoutingPlan.from_orderings( context, comm, input_ordering, output_ordering, output_boundaries ) buffer = _OutputPartitionBuffer(context, ch_in, output_boundaries, output_ordering) @@ -519,8 +545,8 @@ async def _adjust_ordering_impl( plan.local_window[1], ) - # Before receiving remote pieces, emit the local-only prefix and send all - # remote-owned pieces this rank is responsible for. + # Pre-exchange phase: emit the local-only prefix and send remote-owned + # pieces this rank is responsible for. pre_start, pre_stop = plan.local_window[0], first_blocked_pid if plan.owed_remote_pids: pre_start = min(pre_start, plan.owed_remote_range[0]) @@ -545,7 +571,7 @@ async def _adjust_ordering_impl( if plan.local_window[0] <= pid < first_blocked_pid: await _emit_partition( context, - ref_ir, + schema_ir, ir_context, ch_out, pid, @@ -584,9 +610,9 @@ async def _adjust_ordering_impl( TableChunk.from_message(msg, br=context.br()) for msg in pid_store ) chunk = ( - await concat_batch(chunks, context, ref_ir.schema, ir_context) + await concat_batch(chunks, context, schema_ir.schema, ir_context) if chunks - else empty_table_chunk(ref_ir, context, ir_context.get_cuda_stream()) + else empty_table_chunk(schema_ir, context, ir_context.get_cuda_stream()) ) await ch_out.send(context, Message(pid, chunk)) await buffer.assert_input_drained() @@ -596,7 +622,7 @@ async def _adjust_ordering_impl( async def adjust_ordering( context: Context, comm: Communicator, - ref_ir: IR, + schema_ir: IR, ir_context: IRExecutionContext, ch_out: Channel[TableChunk], ch_in: Channel[TableChunk], @@ -606,7 +632,7 @@ async def adjust_ordering( collective_id: int | None = None, ) -> None: """ - Adjust flat Ordering boundaries using contiguous partition ownership. + Adjust Ordering boundaries using contiguous partition ownership. Parameters ---------- @@ -614,7 +640,7 @@ async def adjust_ordering( The streaming context. comm The communicator. - ref_ir + schema_ir An IR node describing the payload schema. ir_context The IR execution context. @@ -635,6 +661,9 @@ async def adjust_ordering( caller is responsible for receiving input metadata and sending output metadata. Input rows are assumed to be globally ordered by ``input_ordering``; sortedness is not checked here. + + The current implementation assumes contiguous partition ownership, strict + output boundaries, and output keys that are a prefix of the input keys. """ _validate_orderings(input_ordering, output_ordering) @@ -645,7 +674,7 @@ async def adjust_ordering( await _adjust_ordering_impl( context, comm, - ref_ir, + schema_ir, ir_context, ch_out, ch_in, From 7150fecd74e27a84a1f3fca63769da711e1bff74 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 15 Jul 2026 07:30:29 -0700 Subject: [PATCH 24/26] require collective id --- .../actor_graph/collectives/ordering.py | 10 +- .../tests/streaming/test_adjust_ordering.py | 102 ++++++++---------- 2 files changed, 46 insertions(+), 66 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 9fde72a867b0..56cac15c21a9 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -513,7 +513,7 @@ async def _adjust_ordering_impl( ch_in: Channel[TableChunk], input_ordering: Ordering, output_ordering: Ordering, - collective_id: int | None, + collective_id: int, ) -> None: """Adjust ordering while using exchange only for remote dependencies.""" output_boundaries = output_ordering.get_boundaries(context.br()) @@ -524,7 +524,6 @@ async def _adjust_ordering_impl( exchange = None if plan.remote_sources or plan.remote_destinations: - assert collective_id is not None exchange = SparseAlltoall( context, comm, @@ -629,7 +628,7 @@ async def adjust_ordering( input_ordering: Ordering, output_ordering: Ordering, *, - collective_id: int | None = None, + collective_id: int, ) -> None: """ Adjust Ordering boundaries using contiguous partition ownership. @@ -662,14 +661,11 @@ async def adjust_ordering( metadata. Input rows are assumed to be globally ordered by ``input_ordering``; sortedness is not checked here. - The current implementation assumes contiguous partition ownership, strict + The current implementation requires contiguous partition ownership, strict output boundaries, and output keys that are a prefix of the input keys. """ _validate_orderings(input_ordering, output_ordering) - if comm.nranks > 1 and collective_id is None: - raise ValueError("collective_id is required when comm.nranks > 1.") - try: await _adjust_ordering_impl( context, diff --git a/python/cudf_polars/tests/streaming/test_adjust_ordering.py b/python/cudf_polars/tests/streaming/test_adjust_ordering.py index c35cd608d24c..31a2015a9b21 100644 --- a/python/cudf_polars/tests/streaming/test_adjust_ordering.py +++ b/python/cudf_polars/tests/streaming/test_adjust_ordering.py @@ -119,7 +119,7 @@ async def _adjust_and_collect( input_ordering: Ordering, output_ordering: Ordering, *, - collective_id: int | None = None, + collective_id: int, ) -> dict[int, pl.DataFrame]: """Run adjustment and collect output chunks by partition ID.""" ch_in = context.create_channel() @@ -190,7 +190,7 @@ async def _adjust_direct( input_ordering: Ordering, output_ordering: Ordering, *, - collective_id: int | None = None, + collective_id: int, ) -> None: ch_in = context.create_channel() ch_out = context.create_channel() @@ -238,29 +238,18 @@ def test_adjust_ordering_rejects_invalid_orderings( stream=stream, ) - with pytest.raises(err, match=match): + with pytest.raises(err, match=match), reserve_op_id() as op_id: asyncio.run( - _adjust_direct(context, spmd_engine.comm, input_ordering, output_ordering) + _adjust_direct( + context, + spmd_engine.comm, + input_ordering, + output_ordering, + collective_id=op_id, + ) ) -@pytest.mark.spmd -def test_adjust_ordering_requires_collective_id( - spmd_engine: SPMDEngine, -) -> None: - context = spmd_engine.context - comm = spmd_engine.comm - if comm.nranks == 1: - pytest.skip("collective_id is only required for multi-rank runs.") - - stream = context.br().stream_pool.get_stream() - input_ordering = _make_ordering(context, 4, stream=stream) - output_ordering = _make_ordering(context, 4, stream=stream) - - with pytest.raises(ValueError, match="collective_id"): - asyncio.run(_adjust_direct(context, comm, input_ordering, output_ordering)) - - @pytest.mark.spmd @pytest.mark.parametrize( "target_boundary,expected", @@ -468,7 +457,7 @@ def test_adjust_ordering_all_empty_input(spmd_engine: SPMDEngine) -> None: if pid * comm.nranks // output_npartitions == comm.rank } - if comm.nranks == 1: + with reserve_op_id() as op_id: output = asyncio.run( _adjust_and_collect( context, @@ -476,20 +465,9 @@ def test_adjust_ordering_all_empty_input(spmd_engine: SPMDEngine) -> None: _frame([]), input_ordering, output_ordering, + collective_id=op_id, ) ) - else: - with reserve_op_id() as op_id: - output = asyncio.run( - _adjust_and_collect( - context, - comm, - _frame([]), - input_ordering, - output_ordering, - collective_id=op_id, - ) - ) _assert_partition_output(output, expected) @@ -502,7 +480,7 @@ def test_adjust_ordering_all_empty_input(spmd_engine: SPMDEngine) -> None: (0, {0: [], 1: list(range(8))}), ], ) -def test_adjust_ordering_single_rank_no_collective( +def test_adjust_ordering_single_rank( spmd_engine: SPMDEngine, target_boundary: int, expected: _ExpectedPartitions, @@ -515,15 +493,17 @@ def test_adjust_ordering_single_rank_no_collective( stream = context.br().stream_pool.get_stream() input_ordering = _make_ordering(context, 4, stream=stream) output_ordering = _make_ordering(context, target_boundary, stream=stream) - output_by_pid = asyncio.run( - _adjust_and_collect( - context, - comm, - _frame(list(range(8))), - input_ordering, - output_ordering, + with reserve_op_id() as op_id: + output_by_pid = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(list(range(8))), + input_ordering, + output_ordering, + collective_id=op_id, + ) ) - ) _assert_partition_output(output_by_pid, expected) @@ -538,15 +518,17 @@ def test_adjust_ordering_multi_chunk_input(spmd_engine: SPMDEngine) -> None: stream = context.br().stream_pool.get_stream() input_ordering = _make_ordering(context, 4, stream=stream) output_ordering = _make_ordering(context, 4, stream=stream) - output = asyncio.run( - _adjust_and_collect( - context, - comm, - [_frame([0, 1]), _frame([2, 3, 4, 5]), _frame([6, 7]), _frame([])], - input_ordering, - output_ordering, + with reserve_op_id() as op_id: + output = asyncio.run( + _adjust_and_collect( + context, + comm, + [_frame([0, 1]), _frame([2, 3, 4, 5]), _frame([6, 7]), _frame([])], + input_ordering, + output_ordering, + collective_id=op_id, + ) ) - ) _assert_partition_output(output, {0: [0, 1, 2, 3], 1: [4, 5, 6, 7]}) @@ -596,14 +578,16 @@ def test_adjust_ordering_respects_order_key_metadata( null_order=null_order, stream=stream, ) - output_by_pid = asyncio.run( - _adjust_and_collect( - context, - comm, - _frame(keys), - input_ordering, - output_ordering, + with reserve_op_id() as op_id: + output_by_pid = asyncio.run( + _adjust_and_collect( + context, + comm, + _frame(keys), + input_ordering, + output_ordering, + collective_id=op_id, + ) ) - ) _assert_partition_output(output_by_pid, expected) From f3bdfede8ef61c78bc85549b1f4365a66711ad55 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 15 Jul 2026 07:51:49 -0700 Subject: [PATCH 25/26] add shutdown_channels_on_error --- .../actor_graph/collectives/ordering.py | 24 +++++-------- .../streaming/actor_graph/utils.py | 34 ++++++++++++++----- 2 files changed, 34 insertions(+), 24 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 56cac15c21a9..6204caa4fabf 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -24,7 +24,7 @@ ChunkStore, concat_batch, empty_table_chunk, - gather_in_task_group, + shutdown_channels_on_error, ) from cudf_polars.utils.cuda_stream import stream_ordered_after @@ -411,20 +411,22 @@ async def assert_input_drained(self) -> None: raise RuntimeError( "adjust_ordering left buffered data after all output was emitted." ) - if ( - not self.input_done - and (msg := await self.ch_in.recv(self.context)) is not None - ): + while not self.input_done: + msg = await self.ch_in.recv(self.context) + if msg is None: + self.input_done = True + break rows = ( TableChunk.from_message(msg, br=self.context.br()) .table_view() .num_rows() ) + if rows == 0: + continue raise RuntimeError( "adjust_ordering left unread input after all output was emitted " f"({rows} rows)" ) - self.input_done = True def _store_chunk( @@ -666,7 +668,7 @@ async def adjust_ordering( """ _validate_orderings(input_ordering, output_ordering) - try: + async with shutdown_channels_on_error(context, ch_in, ch_out): await _adjust_ordering_impl( context, comm, @@ -678,11 +680,3 @@ async def adjust_ordering( output_ordering, collective_id, ) - except BaseException: - await gather_in_task_group( - ch_in.shutdown(context), - ch_in.shutdown_metadata(context), - ch_out.shutdown(context), - ch_out.shutdown_metadata(context), - ) - raise diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index 8e25c8166013..d02077206739 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -221,6 +221,27 @@ async def gather_in_task_group(*coroutines: Coroutine[Any, Any, Any]) -> list[An return [task.result() for task in tasks] +async def shutdown_channels(context: Context, *channels: Channel[Any]) -> None: + """Shutdown data and metadata paths for all channels.""" + await gather_in_task_group( + *itertools.chain.from_iterable( + (ch.shutdown(context), ch.shutdown_metadata(context)) for ch in channels + ) + ) + + +@asynccontextmanager +async def shutdown_channels_on_error( + context: Context, *channels: Channel[Any] +) -> AsyncIterator[None]: + """Shutdown channels on error without actor tracing.""" + try: + yield + except BaseException: + await shutdown_channels(context, *channels) + raise + + @asynccontextmanager async def shutdown_on_error( context: Context, @@ -229,10 +250,10 @@ async def shutdown_on_error( ir_context: IRExecutionContext | None = None, ) -> AsyncIterator[ActorTracer | None]: """ - Shutdown on error for rapidsmpf. + Actor-level shutdown and tracing for rapidsmpf. - This context manager handles channel cleanup on errors and optionally - emits structlog tracing events when LOG_TRACES is enabled. + This context manager handles actor channel cleanup on errors and emits + structlog tracing events. Parameters ---------- @@ -271,12 +292,7 @@ async def shutdown_on_error( try: yield tracer except BaseException: - await gather_in_task_group( - *itertools.chain.from_iterable( - (ch.shutdown(context), ch.shutdown_metadata(context)) - for ch in channels - ) - ) + await shutdown_channels(context, *channels) raise finally: stop = time.monotonic_ns() From 112449b2a8a0e0e98651bb0ff04aa251cebec2c7 Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 15 Jul 2026 08:19:41 -0700 Subject: [PATCH 26/26] avoid pid column --- .../actor_graph/collectives/ordering.py | 112 ++++++++---------- 1 file changed, 49 insertions(+), 63 deletions(-) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py index 6204caa4fabf..805c00c29e92 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/ordering.py @@ -41,7 +41,6 @@ _PID_DTYPE = DataType(pl.Int32()) -_PID_PLC_DTYPE = plc.DataType(plc.TypeId.INT32) _PartitionRange = tuple[int, int] @@ -194,21 +193,6 @@ def _split_points( ) -def _append_partition_id(table: plc.Table, pid: int, stream: Stream) -> plc.Table: - """ - Append the target partition ID before packing a remote table piece. - - SparseAlltoall exchanges packed table payloads, so the partition ID travels - as a temporary column and is stripped by the receiver. - """ - pid_col = plc.Column.from_scalar( - plc.Scalar.from_py(pid, _PID_PLC_DTYPE, stream=stream), - table.num_rows(), - stream=stream, - ) - return plc.Table([*table.columns(), pid_col]) - - def _boundary_search_positions( input_boundary_table: plc.Table, output_boundary_table: plc.Table, @@ -298,27 +282,28 @@ def _sources_for_pid( ] -def _unpack_remote_piece( +def _remote_pids_from_source( + plan: _RoutingPlan, + source_rank: int, + destination_rank: int, +) -> list[int]: + """Return pids sent from one source rank to one destination rank.""" + return [ + pid + for pid in range(*plan.source_ranges[source_rank]) + if _contiguous_owner(pid, len(plan.source_ranges), plan.npartitions) + == destination_rank + ] + + +def _unpack_remote_partition( packed: PackedData, stream: Stream, br: BufferResource, -) -> tuple[int, TableChunk]: - """Unpack one remote piece and recover its temporary target partition ID.""" - table = unpack_and_concat([packed], stream=stream, br=br) - *payload_cols, pid_col = table.columns() - pid = int( - DataFrame.from_table( - plc.Table([pid_col]), - ["pid"], - [_PID_DTYPE], - stream, - ) - .to_polars() - .item(0, 0) - ) - payload = plc.Table(payload_cols).copy(stream=stream, mr=br.device_mr) - return pid, TableChunk.from_pylibcudf_table( - payload, +) -> TableChunk: + """Unpack one remote output-partition payload.""" + return TableChunk.from_pylibcudf_table( + unpack_and_concat([packed], stream=stream, br=br), stream, exclusive_view=True, br=br, @@ -440,41 +425,33 @@ def _store_chunk( stores[pid].insert(Message(pid, chunk)) -async def _send_remote_store( - context: Context, - comm: Communicator, - exchange: SparseAlltoall, - npartitions: int, - pid: int, - store: ChunkStore, -) -> None: - """Send all locally-read pieces for one remote-owned output partition.""" - for msg in store: - await _send_remote_piece( - context, - comm, - exchange, - npartitions, - pid, - TableChunk.from_message(msg, br=context.br()), - ) - - -async def _send_remote_piece( +async def _send_remote_partition( context: Context, comm: Communicator, + schema_ir: IR, + ir_context: IRExecutionContext, exchange: SparseAlltoall, npartitions: int, pid: int, - chunk: TableChunk, + store: ChunkStore | None, ) -> None: - """Send one output-partition piece to its remote owner.""" + """Send one packed payload for one remote-owned output partition.""" + chunks = ( + [TableChunk.from_message(msg, br=context.br()) for msg in store] + if store is not None + else [] + ) + chunk = ( + await concat_batch(chunks, context, schema_ir.schema, ir_context) + if chunks + else empty_table_chunk(schema_ir, context, ir_context.get_cuda_stream()) + ) stream = chunk.stream exchange.insert( _contiguous_owner(pid, comm.nranks, npartitions), packed_data_from_cudf_packed_columns( pack( - _append_partition_id(chunk.table_view(), pid, stream), + chunk.table_view(), stream, mr=context.br().device_mr, ), @@ -555,13 +532,16 @@ async def _adjust_ordering_impl( pieces_by_source: dict[int, dict[int, ChunkStore]] = {comm.rank: {}} local_pieces = pieces_by_source[comm.rank] + owed_remote_pid_set = set(plan.owed_remote_pids) for pid in range(pre_start, pre_stop): piece = await buffer.collect_output_partition(pid) owner = _contiguous_owner(pid, comm.nranks, plan.npartitions) - if exchange is not None and owner != comm.rank and piece is not None: - await _send_remote_store( + if exchange is not None and pid in owed_remote_pid_set: + await _send_remote_partition( context, comm, + schema_ir, + ir_context, exchange, plan.npartitions, pid, @@ -587,9 +567,15 @@ async def _adjust_ordering_impl( for source_rank in plan.remote_sources: remote_pieces: dict[int, ChunkStore] = {} stream = context.br().stream_pool.get_stream() - for packed in exchange.extract(source_rank): - pid, chunk = _unpack_remote_piece(packed, stream, context.br()) - _store_chunk(context, remote_pieces, pid, chunk) + expected_pids = _remote_pids_from_source(plan, source_rank, comm.rank) + for pid, packed in zip( + expected_pids, + exchange.extract(source_rank), + strict=True, + ): + chunk = _unpack_remote_partition(packed, stream, context.br()) + if chunk.table_view().num_rows() > 0: + _store_chunk(context, remote_pieces, pid, chunk) pieces_by_source[source_rank] = remote_pieces # Collect and emit the remaining output partitions for this rank.