diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py index 1c59e207ba5d..9ec03c02881d 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py @@ -386,9 +386,8 @@ def execute_ir_on_rank( stats: StatsCollector, collective_id_map: dict[IR, list[int]], *, - collect_metadata: bool = False, query_id: uuid.UUID, -) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata]]: """ Execute a Polars IR query on a single rank's GPU. @@ -414,8 +413,6 @@ def execute_ir_on_rank( Statistics collector. collective_id_map Mapping from IR nodes to their pre-allocated collective operation IDs. - collect_metadata - Whether to collect channel metadata during execution. query_id Unique identifier for the query, propagated into actor traces. @@ -424,13 +421,12 @@ def execute_ir_on_rank( result This rank's output fragment as a Polars DataFrame. metadata - Collected channel metadata if ``collect_metadata`` is ``True``, - otherwise ``None``. + Collected channel metadata. """ ir_context = IRExecutionContext( get_cuda_stream=ctx.get_stream_from_pool, query_id=query_id ) - metadata_collector: list[ChannelMetadata] | None = [] if collect_metadata else None + metadata_collector: list[ChannelMetadata] = [] nodes, output = generate_network( ctx, @@ -606,9 +602,8 @@ def evaluate_on_rank( ir: IR, config_options: ConfigOptions[StreamingExecutor], *, - collect_metadata: bool = False, query_id: uuid.UUID, -) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]: +) -> tuple[pl.DataFrame, list[ChannelMetadata]]: """ Evaluate a polars IR plan on a single rank. @@ -632,8 +627,6 @@ def evaluate_on_rank( Root of the **pre-lowered** IR graph. config_options Executor configuration forwarded from the client. - collect_metadata - Whether to collect channel metadata during execution. query_id Unique identifier for the query, propagated into actor traces. @@ -642,8 +635,7 @@ def evaluate_on_rank( result This rank's output fragment as a Polars DataFrame. metadata - Collected channel metadata if *collect_metadata* is ``True``, - otherwise ``None``. + Collected channel metadata. """ stats = allgather_stats(comm, ctx.br(), ir, config_options) ir, partition_info = lower_ir_graph(ir, config_options, stats) @@ -663,6 +655,5 @@ def evaluate_on_rank( config_options, stats, collective_id_map, - collect_metadata=collect_metadata, query_id=query_id, ) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py index aee0696f2af4..d168d7a02aaa 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py @@ -411,15 +411,24 @@ def _worker_evaluate( mp_ctx: _WorkerContext = getattr(dask_worker, f"_cudf_polars_mp_context_{uid}") if mp_ctx.ctx is None or mp_ctx.comm is None or mp_ctx.py_executor is None: raise RuntimeError("_setup_worker must be called before _worker_evaluate") - return evaluate_on_rank( + # evaluate_on_rank always collects metadata internally so we can read + # metadata[-1].duplicated to decide whether to suppress this rank's output. + # The client concatenates each rank's result, so without this dedup an + # output marked duplicated=True would appear N times. The external + # collect_metadata parameter still controls whether the collected list is + # returned to the client (see the return statement), which is the cost we + # care about saving when the caller doesn't need the metadata. + df, metadata = evaluate_on_rank( mp_ctx.ctx, mp_ctx.comm, mp_ctx.py_executor, ir, config_options, - collect_metadata=collect_metadata, query_id=query_id, ) + if mp_ctx.comm.rank != 0 and metadata and metadata[-1].duplicated: + df = df.clear() + return df, metadata if collect_metadata else None def evaluate_pipeline_dask_mode( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py index 2ef7412ab19b..f3ccc669c532 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py @@ -368,21 +368,32 @@ def evaluate_polars_ir( RuntimeError If :meth:`setup_worker` has not been called first. """ - if self._ctx is None: + if self._ctx is None or self._comm is None: raise RuntimeError("setup_worker must be called before evaluate_polars_ir") # Ray transfers the returned Polars DataFrame back to the client via the # object store (pickle / Arrow IPC). The DataFrame is already on CPU at # this point (to_polars() copies the result off-GPU), so no GPU memory # crosses process boundaries. - return evaluate_on_rank( + # + # evaluate_on_rank always collects metadata internally so we can read + # metadata[-1].duplicated to decide whether to suppress this rank's + # output. The client concatenates each rank's result, so without this + # dedup an output marked duplicated=True would appear N times. The + # external collect_metadata parameter still controls whether the + # collected list is returned to the client (see the return statement), + # which is the cost we care about saving when the caller doesn't need + # the metadata. + df, metadata = evaluate_on_rank( self._ctx, self._comm, self._py_executor, ir, config_options, - collect_metadata=collect_metadata, query_id=query_id, ) + if self._comm.rank != 0 and metadata and metadata[-1].duplicated: + df = df.clear() + return df, metadata if collect_metadata else None def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> T: return func(*args, **kwargs) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py index e7b0cc936b91..c50a09493a5e 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py @@ -108,15 +108,15 @@ def evaluate_pipeline_spmd_mode( context = config_options.executor.spmd_context.context py_executor = config_options.executor.spmd_context.py_executor - return evaluate_on_rank( + df, metadata = evaluate_on_rank( context, comm, py_executor, ir, config_options, - collect_metadata=collect_metadata, query_id=query_id, ) + return df, metadata if collect_metadata else None def allgather_polars_dataframe( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py index 330e7d208959..7798f1faa5df 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py @@ -207,7 +207,7 @@ async def _collect_small_side_for_broadcast( stream = ir_context.get_cuda_stream() gathered = await allgather.extract_concatenated(stream) # When every rank inserted zero chunks, the AllGather has no schema - # to infer and returns a 0-column table. Substitute a properly typed + # to infer and returns a 0 column table. Substitute a properly typed # empty table for the small side so downstream joins still match the # expected schema. table = ( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py index 56509a93f06a..fd83c5e092f2 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py @@ -732,20 +732,10 @@ async def metadata_drain_node( ): # Drain metadata channel (we don't need it after this point) metadata = await recv_metadata(ch_in, context) - send_empty = metadata.duplicated and comm.rank != 0 if metadata_collector is not None: metadata_collector.append(metadata) - # Forward non-duplicated data messages while (msg := await ch_in.recv(context)) is not None: - if not send_empty: - await ch_out.send(context, msg) - - # Send empty data if needed - if send_empty: - stream = ir_context.get_cuda_stream() - await ch_out.send( - context, Message(0, empty_table_chunk(ir, context, stream)) - ) + await ch_out.send(context, msg) await ch_out.drain(context) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py index 2484620234df..df4c6d5dc949 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py @@ -17,6 +17,7 @@ from cudf_polars.experimental.rapidsmpf.nodes import define_actor, shutdown_on_error from cudf_polars.experimental.rapidsmpf.utils import ( ChannelManager, + empty_table_chunk, gather_in_task_group, process_children, recv_metadata, @@ -49,8 +50,7 @@ async def union_node( context The rapidsmpf context. comm - The communicator. Used to suppress duplicated children's chunks on - non-root ranks so they aren't emitted twice cluster-wide. + The communicator. ir The Union IR node. ir_context @@ -69,15 +69,9 @@ async def union_node( metadata = await gather_in_task_group( *(recv_metadata(ch, context) for ch in chs_in) ) - # When a child has duplicated=True, every rank has produced the same - # data and only rank 0 should forward it -- otherwise the downstream - # client-side concat would over-count by `nranks - 1` for each - # duplicated chunk. - skip = tuple(meta.duplicated and comm.rank != 0 for meta in metadata) - total_local_count = sum( - 0 if drop else meta.local_count - for meta, drop in zip(metadata, skip, strict=True) - ) + # Chunk counts on the wire are uniform across ranks, so report the + # full sum. + total_local_count = sum(meta.local_count for meta in metadata) duplicated = all(meta.duplicated for meta in metadata) await send_metadata( ch_out, @@ -88,23 +82,29 @@ async def union_node( ), ) + # When a child has duplicated=True, every rank has produced the same + # rows, so we drop them everywhere except rank 0 to avoid N counting. + suppress = tuple(meta.duplicated and comm.rank != 0 for meta in metadata) + seq_num_offset = 0 - for ch_in, drop in zip(chs_in, skip, strict=True): + for ch_in, drop in zip(chs_in, suppress, strict=True): num_ch_chunks = 0 while (msg := await ch_in.recv(context)) is not None: - if not drop: - await ch_out.send( - context, - Message( - msg.sequence_number + seq_num_offset, - TableChunk.from_message( - msg, br=context.br() - ).make_available_and_spill( - context.br(), allow_overbooking=True - ), - ), - ) - num_ch_chunks += 1 + if drop: + stream = ir_context.get_cuda_stream() + out_chunk = empty_table_chunk(ir, context, stream) + else: + out_chunk = TableChunk.from_message( + msg, br=context.br() + ).make_available_and_spill(context.br(), allow_overbooking=True) + await ch_out.send( + context, + Message( + msg.sequence_number + seq_num_offset, + out_chunk, + ), + ) + num_ch_chunks += 1 seq_num_offset += num_ch_chunks await ch_out.drain(context) diff --git a/python/cudf_polars/cudf_polars/experimental/select.py b/python/cudf_polars/cudf_polars/experimental/select.py index 9ab30f9be136..606741d587d1 100644 --- a/python/cudf_polars/cudf_polars/experimental/select.py +++ b/python/cudf_polars/cudf_polars/experimental/select.py @@ -432,8 +432,9 @@ def _( named_expr = expr.NamedExpr(ir.exprs[0].name or "len", lit_expr) # Use Empty as the input so the streaming network's metadata flows - # `duplicated=True` end-to-end. Without that, every rank emits the - # literal once and the client concatenates N copies. + # duplicated=True end to end. Without that, the literal expression + # would be evaluated chunkwise over child and emit one row per input + # chunk instead of a single row globally. input_ir: IR = Empty({}) new_node = Select( {named_expr.name: named_expr.value.dtype}, diff --git a/python/cudf_polars/tests/experimental/test_spmd.py b/python/cudf_polars/tests/experimental/test_spmd.py index 96ec5eab9320..fabaeedbc780 100644 --- a/python/cudf_polars/tests/experimental/test_spmd.py +++ b/python/cudf_polars/tests/experimental/test_spmd.py @@ -22,6 +22,7 @@ SPMDEngine, allgather_polars_dataframe, ) +from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.utils.config import MemoryResourceConfig if TYPE_CHECKING: @@ -294,6 +295,24 @@ def test_run(spmd_engine: SPMDEngine) -> None: assert result == [os.getpid()] +def test_sort_slice_over_union_of_duplicated_streams( + spmd_engine: SPMDEngine, +) -> None: + """Sort+head over a concat of two group-by branches returns the global result on every rank.""" + lf1 = ( + pl.LazyFrame({"name": ["alice"], "score": [1.0]}) + .group_by("name") + .agg(pl.col("score").sum()) + ) + lf2 = ( + pl.LazyFrame({"name": ["bob"], "score": [2.0]}) + .group_by("name") + .agg(pl.col("score").sum()) + ) + lf = pl.concat([lf1, lf2]).sort("score").head(10) + assert_gpu_result_equal(lf, engine=spmd_engine, check_row_order=False) + + def test_reset_keeps_comm_alive(comm: Communicator) -> None: """``_reset`` must not rebuild the communicator.""" with SPMDEngine( diff --git a/python/cudf_polars/tests/experimental/test_union.py b/python/cudf_polars/tests/experimental/test_union.py index 79a3ca649632..ecb8f1ebfe19 100644 --- a/python/cudf_polars/tests/experimental/test_union.py +++ b/python/cudf_polars/tests/experimental/test_union.py @@ -17,3 +17,19 @@ def test_union_shared_fanout_no_deadlock(streaming_engine): project = df.select("key", "val") q = pl.concat([gb, project]) assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) + + +def test_sort_slice_over_union_of_duplicated_streams(streaming_engine): + # Sort+head over a concat of two group-by branches. + lf1 = ( + pl.LazyFrame({"name": ["alice"], "score": [1.0]}) + .group_by("name") + .agg(pl.col("score").sum()) + ) + lf2 = ( + pl.LazyFrame({"name": ["bob"], "score": [2.0]}) + .group_by("name") + .agg(pl.col("score").sum()) + ) + q = pl.concat([lf1, lf2]).sort("score").head(10) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) diff --git a/python/cudf_polars/tests/test_groupby.py b/python/cudf_polars/tests/test_groupby.py index c649a0997c50..728cc09e76a0 100644 --- a/python/cudf_polars/tests/test_groupby.py +++ b/python/cudf_polars/tests/test_groupby.py @@ -484,13 +484,8 @@ def test_groupby_sum_decimal_null_group(engine: pl.GPUEngine) -> None: @pytest.mark.xfail( - raises=(AssertionError, pl.exceptions.SchemaError), - reason=( - "https://github.com/rapidsai/cudf/issues/19610 — in-memory engine " - "fails with AssertionError (wrong values); multi-rank streaming " - "fails earlier with SchemaError (literal agg yields a divergent " - "schema after cross-rank concat)." - ), + raises=AssertionError, + reason="https://github.com/rapidsai/cudf/issues/19610", ) def test_groupby_literal_agg(engine: pl.GPUEngine): df = pl.LazyFrame({"c0": [True, False]})