Add remaining cudf_streaming tests and benchmarks - #22814
Merged
rapids-bot[bot] merged 28 commits intoJun 9, 2026
Conversation
github-actions
Bot
added
libcudf
vyasr
force-pushed
the
feat/add_remaining_rapidsmpf_pieces
branch
from
June 9, 2026 05:20
207dc37 to
f1efac5
Compare
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds BUILD_BENCHMARKS/BUILD_EXAMPLES, many C++ benchmark targets and NDSh streaming operators, shared benchmark utilities, Python validation tooling and examples, plus extensive C++ and Python tests. ChangesStreaming benchmarks and utilities
ND-shuffle streaming operators and query benches
Examples and validation tooling
Streaming and shuffler tests
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Note
Due to the large number of review comments, Critical severity comments were prioritized as inline comments.
🟠 Major comments (18)
cpp/libcudf_streaming/tests/streaming/test_leaf_actor.cpp-147-148 (1)
147-148:⚠️ Potential issue | 🟠 Major | ⚡ Quick winInitialize atomic accumulators before concurrent use.
Both
resultatomics are default-initialized and then potentially mutated viafetch_add; this is undefined behavior and can make these tests flaky.Suggested fix
- std::atomic<int> result; + std::atomic<int> result{0}; ... - std::atomic<int> result; + std::atomic<int> result{0};Also applies to: 164-166
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/tests/streaming/test_leaf_actor.cpp` around lines 147 - 148, The atomic accumulators (e.g., the local std::atomic<int> variables named result that are passed into consumer(ctx, ch, result)) must be value-initialized before any concurrent use; change their declaration to initialize them to 0 (e.g., std::atomic<int> result{0};) wherever you declare them (both occurrences used with consumer and the other consumer call around those lines) so fetch_add operates on a defined value.cpp/libcudf_streaming/tests/test_shuffler.cpp-327-329 (1)
327-329:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUse
future::get()so worker failures are not hidden.
wait()only blocks; it does not propagate exceptions from async tasks, so this can silently pass when worker threads fail.Suggested fix
- for (auto& f : futures) { - ASSERT_NO_THROW(f.wait()); - } + for (auto& f : futures) { + ASSERT_NO_THROW(f.get()); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/tests/test_shuffler.cpp` around lines 327 - 329, The loop over the futures currently calls f.wait(), which only blocks and hides exceptions from async workers; replace the wait() call in the for (auto& f : futures) loop with f.get() and keep the ASSERT_NO_THROW around that call so any exceptions thrown by the worker tasks (from the futures vector) are propagated and cause the test to fail as intended.python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py-366-369 (1)
366-369:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH:
--ray-addressis accepted but ignoredLine 367 always uses
address="auto"even whenargs.ray_addressis provided, so explicit cluster targeting is not honored.Suggested fix
- if args.ray_address or os.environ.get("RAY_ADDRESS") is not None: - ray.init(address="auto") # connect to existing cluster + if args.ray_address is not None: + ray.init(address=args.ray_address) + elif os.environ.get("RAY_ADDRESS") is not None: + ray.init(address="auto") # connect to existing cluster else: ray.init(num_gpus=args.num_workers, dashboard_host="0.0.0.0")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py` around lines 366 - 369, The code always calls ray.init(address="auto") even when args.ray_address is supplied; update the connection branch in the ray.init call to use the actual provided address. In the block that checks args.ray_address or os.environ.get("RAY_ADDRESS"), compute an address value like address = args.ray_address or os.environ.get("RAY_ADDRESS") (falling back to "auto" only if neither is set) and pass that into ray.init(address=address) so explicit --ray-address is honored; leave the else branch (ray.init(num_gpus=..., dashboard_host=...)) unchanged.python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py-372-372 (1)
372-372:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH: Input path collection can feed invalid files into Parquet reader
Line 372 recursively collects all filesystem entries, including directories and non-Parquet files.
Suggested fix
- paths=sorted(map(str, args.input.glob("**/*"))), + paths=sorted( + str(p) for p in args.input.rglob("*.parquet") if p.is_file() + ),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py` at line 372, The current paths=sorted(map(str, args.input.glob("**/*"))) collects directories and non-Parquet files which can break the Parquet reader; change the collection to only regular files with Parquet extensions (e.g., use args.input.rglob("**/*.parquet") or filter args.input.rglob("*") by p.is_file() and p.suffix in {".parquet", ".parquet.gzip", ".parquet.gz"}), update the variable referenced as paths so only valid file paths (strings) are produced before passing to the Parquet reader.python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py-391-391 (1)
391-391:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH: Input file discovery includes directories/non-Parquet paths
Line 391 uses
glob("**/*"), which can pass directories or unrelated files toread_parquet, causing runtime failures.Suggested fix
- paths=sorted(map(str, args.input.glob("**/*"))), + paths=sorted( + str(p) for p in args.input.rglob("*.parquet") if p.is_file() + ),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py` at line 391, Input discovery uses args.input.glob("**/*") which yields directories and non-Parquet files and can break read_parquet; change the discovery to only include regular files with Parquet extensions. Update the code that builds paths (the variable assigned with paths=sorted(...)) to either use args.input.glob("**/*.parquet") (and include ".pq" if needed) or filter the globbed entries by Path.is_file() and by suffix in {".parquet",".pq"} before converting to str, so downstream calls to read_parquet only receive valid file paths.python/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.py-76-80 (1)
76-80:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH: Odd
num_rowsproduces mismatched column lengthsLine 77 creates only
2 * (num_rows // 2)strings. With oddnum_rows, table construction fails because column lengths differ.Suggested fix
plc.Column.from_iterable_of_py( - ["cat", "dog"] * (self._num_rows // 2), + (["cat", "dog"] * ((self._num_rows + 1) // 2))[ + : self._num_rows + ], plc.DataType(plc.TypeId.STRING), ),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.py` around lines 76 - 80, The string column generation uses ["cat", "dog"] * (self._num_rows // 2) which produces only 2*(self._num_rows//2) entries and breaks for odd self._num_rows; change the construction used in plc.Column.from_iterable_of_py so it always yields exactly self._num_rows items (e.g., generate values with a bounded iterator or list comprehension that repeats "cat"/"dog" for range(self._num_rows) and slices/truncates as needed) so the column length matches self._num_rows.python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py-205-211 (1)
205-211:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH: Configured spill limits are bypassed in the shuffle path
Line 205 recreates
brfrom the current device resource, which discards the caller-providedBufferResource(including--spill-devicelimits). This can invalidate intended memory pressure behavior and increase OOM risk.Suggested fix
- br = BufferResource(rmm.mr.get_current_device_resource()) shuffler = Shuffler( comm, op_id=0, total_num_partitions=total_num_partitions, br=br, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py` around lines 205 - 211, The code is recreating a new BufferResource via BufferResource(rmm.mr.get_current_device_resource()) which discards the caller-provided BufferResource (and its --spill-device limits); instead, stop reconstructing br and pass the original caller-provided BufferResource into the Shuffler. Locate the BufferResource creation around the Shuffler construction (symbol names: BufferResource, br, Shuffler) and remove or replace the BufferResource(...) instantiation so that the existing br instance with configured spill limits is used when constructing Shuffler.python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py-209-213 (1)
209-213:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH:
read_and_insertcrashes when a worker gets zero filesWhen
pathsis empty (possible whennum_workers > num_input_files),column_namesis never assigned and Line 213 raisesUnboundLocalError.Suggested fix
def read_and_insert(self, paths: list[str]) -> list[str]: @@ - for i in range(0, len(paths), self.batchsize): + column_names: list[str] = [] + for i in range(0, len(paths), self.batchsize): tbl, column_names = self.read_batch(paths[i : i + self.batchsize]) self.insert_chunk(tbl, column_names) self.insert_finished() return column_names🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py` around lines 209 - 213, The function read_and_insert can raise UnboundLocalError because column_names is only set inside the for-loop when read_batch is called; if paths is empty the loop never runs and the return on column_names fails. Fix by initializing column_names before the loop (e.g., column_names = [] or None) and ensure the function still calls insert_finished() and returns that safe default when no batches were processed; update references in read_and_insert to use this initialized value so insert_chunk, insert_finished, and the final return work even when num_workers > num_input_files.python/cudf_streaming/cudf_streaming/examples/streaming_basic_example.py-45-49 (1)
45-49:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH: Ensure
Context.shutdown()always executes on exception paths.
ctx.shutdown()only runs on the success path. If actor execution/assertion fails, cleanup is skipped and resources can leak.Proposed fix
- # Run all actors. This blocks until every actor has completed. - run_actor_network( - ctx, - actors=( - actor1, - actor2, - actor3, - ), - ) - - # Collect and verify results. - expect = 0 - for msg in out_messages.release(): - table = TableChunk.from_message(msg, br=ctx.br()).table_view() - expect += table.num_rows() - assert total_num_rows[0] == expect - - # Shut down the context explicitly to ensure it happens on the same thread that - # created it. Alternatively, use `with Context(...) as ctx:` to shut it down - # automatically. - ctx.shutdown() - - return total_num_rows[0] + try: + # Run all actors. This blocks until every actor has completed. + run_actor_network( + ctx, + actors=( + actor1, + actor2, + actor3, + ), + ) + + # Collect and verify results. + expect = 0 + for msg in out_messages.release(): + table = TableChunk.from_message(msg, br=ctx.br()).table_view() + expect += table.num_rows() + assert total_num_rows[0] == expect + return total_num_rows[0] + finally: + # Keep cleanup on creator thread, even on failure. + ctx.shutdown()As per coding guidelines, Python changes should handle GPU/resource cleanup safely when exceptions occur (
python/**/*.{py,pyx}: proper cleanup and leak prevention).Also applies to: 129-150
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudf_streaming/cudf_streaming/examples/streaming_basic_example.py` around lines 45 - 49, The Context instance created as ctx may not be shut down on exception; ensure Context.shutdown() always runs by wrapping the actor logic that uses ctx (the block after ctx = Context(...)) in a try/finally (or convert Context to a contextmanager) so that ctx.shutdown() is called in the finally block; locate references to ctx and Context.shutdown() to update both the main usage and the similar section around lines 129-150 so cleanup runs on both success and exception paths.Source: Coding guidelines
cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hpp-13-15 (1)
13-15:⚠️ Potential issue | 🟠 MajorHIGH: Make
parquet_writer.hppself-contained by including<string>
cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hppusesstd::vector<std::string>inwrite_parquet(...)but does not include<string>, which can break builds depending on include order.Suggested fix
`#include` <memory> +#include <string> `#include` <vector>🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hpp` around lines 13 - 15, The header parquet_writer.hpp is not self-contained because write_parquet(...) uses std::vector<std::string> but the file does not include <string>; add `#include` <string> to the header so it explicitly declares std::string and guarantees compilation regardless of include order, ensuring the declaration of write_parquet (and any other uses of std::string in that header) compiles correctly.cpp/libcudf_streaming/benchmarks/bench_pack.cpp-19-20 (1)
19-20: 🛠️ Refactor suggestion | 🟠 Major | 🏗️ Heavy liftHIGH: Benchmark target is implemented with Google Benchmark instead of NVBench
This
bench_*file registers Google Benchmark tests (BENCHMARK(...)), which conflicts with the repository benchmark policy for C++ benchmark files.As per coding guidelines:
cpp/**/*bench*.{cu,cpp}must use NVBench, not Google Benchmark.Also applies to: 228-239, 295-303
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/bench_pack.cpp` around lines 19 - 20, The file uses Google Benchmark (include <benchmark/benchmark.h> and BENCHMARK(...) registrations) but must use NVBench per project policy; replace the Google Benchmark include and BENCHMARK(...) usages with NVBench equivalents (include nvbench/nvbench.cuh or nvbench/nvbench.hpp), convert each BENCHMARK-registered function to an nvbench-style benchmark callback that accepts nvbench::state (or nvbench::benchmark&) and register with NVBENCH_BENCH or nvbench::register_benchmark, and update any benchmark-specific API calls accordingly (search for the symbol BENCHMARK and the include <benchmark/benchmark.h>, plus the benchmark functions referenced around the sections noted such as the blocks at 228-239 and 295-303) so the file compiles and conforms to NVBench conventions.Source: Coding guidelines
cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp-404-456 (1)
404-456:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH:
left_semi_join_shuffleexits without drainingch_outThis actor can terminate with pending output work not fully flushed, unlike the other join actors in this file.
Suggested fix
streaming::Actor left_semi_join_shuffle(...){ @@ while (!ch_out->is_shutdown()) { @@ co_await ch_out->send(semi_join_chunk(...)); } + co_await ch_out->drain(ctx->executor()); }cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp-539-565 (1)
539-565:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH:
-r 0can leaveelapsed_vecempty beforeharmonic_mean
elapsed_vecis only populated for measured runs, sonum_runs == 0can trigger undefined mean/throughput reporting and potential failure.Suggested fix
@@ } + RAPIDSMPF_EXPECTS(num_runs > 0, "-r <num> must be >= 1"); @@ { auto const elapsed_mean = harmonic_mean(elapsed_vec);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp` around lines 539 - 565, The loop may leave elapsed_vec empty when args.num_runs == 0 (e.g., -r 0), causing harmonic_mean(elapsed_vec) to be called on an empty container; update the post-loop block that computes elapsed_mean to first check elapsed_vec.empty() and handle that case (e.g., log a clear message via log->print and skip/short-circuit reporting or set a safe sentinel like elapsed_mean = 0/NaN) instead of calling harmonic_mean; modify the code around the harmonic_mean call (the block that builds the "means: ..." ss) to branch on elapsed_vec.empty() and only compute/format throughput when there are measured runs.cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp-441-608 (1)
441-608:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH:
mainhas no top-level exception boundary for multi-rank executionAn exception after initialization can terminate one rank while peers remain blocked on collectives/barriers.
Suggested fix
int main(int argc, char** argv) { + try { bool use_bootstrap = rapidsmpf::bootstrap::is_running_with_rrun(); @@ if (!use_bootstrap) { RAPIDSMPF_MPI(MPI_Finalize()); } return 0; + } catch (std::exception const& e) { + std::cerr << "Fatal benchmark error: " << e.what() << std::endl; + if (rapidsmpf::mpi::is_initialized()) { RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); } + return 1; + } }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp` around lines 441 - 608, main lacks a top-level exception boundary, so an exception can kill one rank while others remain blocked; wrap the primary body of main (everything after MPI_Init_thread / before MPI_Finalize) in a try/catch that catches std::exception and (...) and on error logs the exception via comm->logger() (or std::cerr if comm isn't constructed), ensures cooperative shutdown by calling comm-related abort/finalize (e.g., MPI_Abort through RAPIDSMPF_MPI when !use_bootstrap or comm-specific abort method if available), and then returns non-zero; ensure MPI_Finalize is still called or MPI_Abort invoked in the catch to avoid hanging ranks.Source: Linters/SAST tools
cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp-287-287 (1)
287-287:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftHIGH: Single-rank constraint contradicts multi-rank initialization
Issue: Code initializes MPI with
MPI_THREAD_MULTIPLEand UCX communicators that support multiple ranks, then requires exactly 1 rank
Why: Makes the benchmark unusable for multi-rank shuffle testing despite all infrastructure being presentConsider removing this check or updating the error message to clarify that multi-rank support is planned but not yet implemented.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp` at line 287, The RAPIDSMPF_EXPECTS check forcing comm->nranks() == 1 conflicts with earlier MPI_THREAD_MULTIPLE and UCX multi-rank setup; either remove that check to enable multi-rank runs or, if multi-rank behavior isn't implemented yet, change the RAPIDSMPF_EXPECTS message to explicitly state "multi-rank execution not yet supported" so it doesn't mislead. Locate the assertion using the symbol RAPIDSMPF_EXPECTS and the comm->nranks() call in bench_streaming_shuffle.cpp and either delete the check or replace the error string to clarify the current single-rank limitation.cpp/libcudf_streaming/benchmarks/utils/misc.hpp-55-71 (1)
55-71:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMEDIUM: Potential integer overflow in parse_integer
Issue: Range validation uses
int64_tbounds, butstatic_cast<T>(val)at line 70 may overflow ifTis a smaller integer type
Why: Defaultmax_valisstd::numeric_limits<std::int64_t>::max(), which exceeds the range of smaller types likeint32_tExample: Parsing "5000000000" into
int32_twould pass validation (< INT64_MAX) but overflow the cast.🛡️ Suggested fix
template <typename T> void parse_integer(T& output, std::string const& str, std::int64_t min_val = 0, - std::int64_t max_val = std::numeric_limits<std::int64_t>::max()) + std::int64_t max_val = std::numeric_limits<T>::max()) { long long val; try { val = std::stoll(str); } catch (std::invalid_argument const&) { RAPIDSMPF_FAIL("cannot parse \"" + str + "\"", std::invalid_argument); } catch (std::out_of_range const&) { RAPIDSMPF_FAIL("\"" + str + "\" is out of range", std::out_of_range); } + RAPIDSMPF_EXPECTS(val >= std::numeric_limits<T>::min() && val <= std::numeric_limits<T>::max(), + "\"" + str + "\" is out of range for target type"); RAPIDSMPF_EXPECTS(min_val <= val && val <= max_val, "\"" + str + "\" is out of range"); output = static_cast<T>(val); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/utils/misc.hpp` around lines 55 - 71, parse_integer currently validates the parsed long long (val) only against min_val/max_val typed as int64_t then casts to T, which can overflow for smaller integer types; update parse_integer to compute the effective bounds as the intersection of the provided min_val/max_val and std::numeric_limits<T>::min()/max(), validate val against that intersection before assigning to output, and produce the same error messages when out-of-range; reference the function parse_integer, parameters output, str, min_val, max_val, and the local val when making this change.cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp-100-116 (1)
100-116:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHIGH: Replace raw
newowner allocations with RAII in filter builders.
owneris heap-allocated via rawnewand only wrapped at the end. If any intermediate allocation throws, the allocation leaks (same pattern in both helpers).Proposed fix
template <typename timestamp_type> std::unique_ptr<streaming::Filter> make_date_filter(rmm::cuda_stream_view stream, cuda::std::chrono::year_month_day date, std::string const& column_name, cudf::ast::ast_operator op) { - auto owner = new std::vector<std::any>; + auto owner = std::make_unique<std::vector<std::any>>(); auto sys_days = cuda::std::chrono::sys_days(date); owner->push_back(std::make_shared<cudf::timestamp_scalar<timestamp_type>>( sys_days.time_since_epoch(), true, stream)); owner->push_back(std::make_shared<cudf::ast::literal>( *std::any_cast<std::shared_ptr<cudf::timestamp_scalar<timestamp_type>>>(owner->at(0)))); owner->push_back(std::make_shared<cudf::ast::column_name_reference>(column_name)); owner->push_back(std::make_shared<cudf::ast::operation>( op, *std::any_cast<std::shared_ptr<cudf::ast::column_name_reference>>(owner->at(2)), *std::any_cast<std::shared_ptr<cudf::ast::literal>>(owner->at(1)))); + auto* owner_raw = owner.release(); return std::make_unique<streaming::Filter>( stream, - *std::any_cast<std::shared_ptr<cudf::ast::operation>>(owner->back()), - OwningWrapper(static_cast<void*>(owner), + *std::any_cast<std::shared_ptr<cudf::ast::operation>>(owner_raw->back()), + OwningWrapper(static_cast<void*>(owner_raw), [](void* p) { delete static_cast<std::vector<std::any>*>(p); })); } template <typename timestamp_type> std::unique_ptr<streaming::Filter> make_date_range_filter( rmm::cuda_stream_view stream, cuda::std::chrono::year_month_day start_date, cuda::std::chrono::year_month_day end_date, std::string const& column_name) { - auto owner = new std::vector<std::any>; + auto owner = std::make_unique<std::vector<std::any>>(); ... + auto* owner_raw = owner.release(); return std::make_unique<streaming::Filter>( stream, - *std::any_cast<std::shared_ptr<cudf::ast::operation>>(owner->back()), - OwningWrapper(static_cast<void*>(owner), + *std::any_cast<std::shared_ptr<cudf::ast::operation>>(owner_raw->back()), + OwningWrapper(static_cast<void*>(owner_raw), [](void* p) { delete static_cast<std::vector<std::any>*>(p); })); }As per coding guidelines, C++ code should avoid raw owning pointers and ensure cleanup in exception paths.
Also applies to: 140-180
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp` around lines 100 - 116, The code heap-allocates owner with raw new in the filter builder (variable owner) and only wraps it late with OwningWrapper, risking leaks if intermediate allocations throw; replace the raw new by creating a local std::vector<std::any> (or a std::shared_ptr<std::vector<std::any>>) on the stack or as a smart pointer, populate it, then pass a managed pointer (e.g., shared_ptr.get() with a capture or move a shared_ptr into OwningWrapper) to streaming::Filter so cleanup happens on exceptions; update both helper sites (the block creating timestamp_scalar, literal, column_name_reference, operation and the return creating OwningWrapper) to use RAII rather than raw new.Sources: Coding guidelines, Linters/SAST tools
cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp-29-32 (1)
29-32:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftHIGH:
write_parquetcannot handle a valid empty result set.The writer infers schema from the first chunk, then throws if
ch_incloses immediately. Any query that legitimately produces zero rows will abort on rank 0 instead of completing cleanly, which makes this output path unusable for empty-result benchmarks/validation runs.Consider passing the expected schema/metadata into
write_parquetso the writer can be initialized before the firstreceive()and the empty-channel case can return without error.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp` around lines 29 - 32, The write_parquet function currently infers schema from the first chunk (using ch_in->receive() and streaming::TableChunk) and errors on an immediately-empty channel; update write_parquet to accept an optional expected schema/metadata parameter (or a prebuilt cudf::table_schema-like object) and use that to initialize the chunked_parquet_writer_options::builder(sink) and underlying writer before calling ch_in->receive(); then handle the case where ch_in closes with no messages by returning success instead of throwing. Locate symbols write_parquet, ch_in->receive(), streaming::TableChunk, and chunked_parquet_writer_options::builder(sink) to implement: add the new parameter, branch to initialize writer from the provided schema, and ensure an empty msg from receive results in a clean return rather than RAPIDSMPF_EXPECTS failure.
🧹 Nitpick comments (1)
cpp/libcudf_streaming/scripts/ndsh.py (1)
446-448: ⚡ Quick winMEDIUM: Prefer strict=True for safety
Issue:
zip(..., strict=False)disables length mismatch checking
Why: If column counts differ unexpectedly, iteration silently stops early, potentially missing validation errors♻️ Suggested fix
for name, out_col, expected_col in zip( - output.column_names, output.columns, expected.columns, strict=False + output.column_names, output.columns, expected.columns, strict=True ):Since schema equality is already verified at line 412, the lengths should always match, making
strict=Truea safer default.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/scripts/ndsh.py` around lines 446 - 448, The zip call over output.column_names, output.columns, and expected.columns currently uses strict=False which can silently ignore length mismatches; change the zip to use strict=True so mismatched column counts raise immediately (since schema equality is already checked at the earlier verification around line 412). Update the call referencing output.column_names, output.columns, and expected.columns in ndsh.py to pass strict=True to zip to enforce length equality during iteration.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/libcudf_streaming/benchmarks/bench_partition.cpp`:
- Around line 49-51: Wrap unchecked CUDA runtime calls in the benchmark with the
CUDF error-checking macro: include the header cudf/utilities/error.hpp and
replace direct calls to cudaGetDeviceProperties(...) and
cudaStreamSynchronize(stream) with CUDF_CUDA_TRY(cudaGetDeviceProperties(...))
and CUDF_CUDA_TRY(cudaStreamSynchronize(stream)); update every occurrence (both
uses of cudaGetDeviceProperties and all cudaStreamSynchronize calls in
bench_partition.cpp) so failures are propagated and do not silently affect
benchmark results.
In `@cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp`:
- Line 349: The call currently does an unchecked
std::dynamic_pointer_cast<rapidsmpf::ucxx::UCXX>(comm)->barrier(); — assign the
cast result to a local std::shared_ptr<rapidsmpf::ucxx::UCXX> (e.g., auto ucxx =
std::dynamic_pointer_cast<rapidsmpf::ucxx::UCXX>(comm)), check that ucxx is not
null, and handle the failure (log and throw std::runtime_error or call
abort/assert) before calling ucxx->barrier(); this prevents dereferencing a null
pointer when the cast fails.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cpp`:
- Around line 83-89: The loop in chunkwise_groupby_requests() (and similarly in
final_groupby_requests()) moves the shared vector aggs into requests via
requests.emplace_back(idx, std::move(aggs)) and then reuses/mutates aggs on the
next iteration, leaving moved-from state and causing incorrect/empty
aggregations; fix by not moving the original aggs — instead create a fresh local
copy (e.g., auto aggs_copy = aggs) inside each loop iteration and
emplace_back(idx, std::move(aggs_copy)) so the original aggs remains intact for
subsequent iterations; apply the same pattern change to final_groupby_requests()
where aggs is currently moved then reused.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp`:
- Around line 339-342: The cudaFree(nullptr) call can return an error which must
be checked; update the startup sequence around cudaFree(nullptr) (near
rapidsmpf::ndsh::FinalizeMPI finalize{} and before cudf::initialize()) to
capture the cudaError_t result and handle non-success returns by logging the
error (including cudaGetErrorString(err)) and terminating/propagating the
failure (or using your project’s CUDA error macro such as CUDA_SAFE_CALL) so
CUDA startup/device-state failures aren’t hidden behind later benchmark errors.
- Around line 237-256: The code computes a sort permutation (indices from
sorted_order) but only materializes the "values" columns into partials, so later
calling cudf::merge(views, keys, ...) uses the original keys positions and
yields wrong ordering; fix by applying the same gather (use the same indices /
split(indices->view(), {k}).front()) to the key columns as well and store
partial tables that contain keys followed by values (or otherwise ensure the
merged tables have key columns at the positions referenced by keys), then pass
those views into cudf::merge; update the construction that pushes into partials
(where cudf::gather is used) to gather chunk.table_view().select(keys) with the
same indices and combine with the gathered values so merge sees remapped keys.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp`:
- Around line 229-243: The latch is being decremented for every received chunk
in filter_grouped_greater (latch->count_down() inside the message loop), which
is wrong because the latch was created with count 1; either move the single
latch->count_down() out of the per-message loop to execute once after all chunks
are processed, or initialize the latch with the total expected chunk count and
keep decrements per chunk; locate the loop that awaits ch_in->receive() and
adjust either the latch construction or the placement of latch->count_down() so
the number of count_down() calls matches the latch initial count.
---
Major comments:
In `@cpp/libcudf_streaming/benchmarks/bench_pack.cpp`:
- Around line 19-20: The file uses Google Benchmark (include
<benchmark/benchmark.h> and BENCHMARK(...) registrations) but must use NVBench
per project policy; replace the Google Benchmark include and BENCHMARK(...)
usages with NVBench equivalents (include nvbench/nvbench.cuh or
nvbench/nvbench.hpp), convert each BENCHMARK-registered function to an
nvbench-style benchmark callback that accepts nvbench::state (or
nvbench::benchmark&) and register with NVBENCH_BENCH or
nvbench::register_benchmark, and update any benchmark-specific API calls
accordingly (search for the symbol BENCHMARK and the include
<benchmark/benchmark.h>, plus the benchmark functions referenced around the
sections noted such as the blocks at 228-239 and 295-303) so the file compiles
and conforms to NVBench conventions.
In `@cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp`:
- Around line 539-565: The loop may leave elapsed_vec empty when args.num_runs
== 0 (e.g., -r 0), causing harmonic_mean(elapsed_vec) to be called on an empty
container; update the post-loop block that computes elapsed_mean to first check
elapsed_vec.empty() and handle that case (e.g., log a clear message via
log->print and skip/short-circuit reporting or set a safe sentinel like
elapsed_mean = 0/NaN) instead of calling harmonic_mean; modify the code around
the harmonic_mean call (the block that builds the "means: ..." ss) to branch on
elapsed_vec.empty() and only compute/format throughput when there are measured
runs.
- Around line 441-608: main lacks a top-level exception boundary, so an
exception can kill one rank while others remain blocked; wrap the primary body
of main (everything after MPI_Init_thread / before MPI_Finalize) in a try/catch
that catches std::exception and (...) and on error logs the exception via
comm->logger() (or std::cerr if comm isn't constructed), ensures cooperative
shutdown by calling comm-related abort/finalize (e.g., MPI_Abort through
RAPIDSMPF_MPI when !use_bootstrap or comm-specific abort method if available),
and then returns non-zero; ensure MPI_Finalize is still called or MPI_Abort
invoked in the catch to avoid hanging ranks.
In `@cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp`:
- Line 287: The RAPIDSMPF_EXPECTS check forcing comm->nranks() == 1 conflicts
with earlier MPI_THREAD_MULTIPLE and UCX multi-rank setup; either remove that
check to enable multi-rank runs or, if multi-rank behavior isn't implemented
yet, change the RAPIDSMPF_EXPECTS message to explicitly state "multi-rank
execution not yet supported" so it doesn't mislead. Locate the assertion using
the symbol RAPIDSMPF_EXPECTS and the comm->nranks() call in
bench_streaming_shuffle.cpp and either delete the check or replace the error
string to clarify the current single-rank limitation.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp`:
- Around line 29-32: The write_parquet function currently infers schema from the
first chunk (using ch_in->receive() and streaming::TableChunk) and errors on an
immediately-empty channel; update write_parquet to accept an optional expected
schema/metadata parameter (or a prebuilt cudf::table_schema-like object) and use
that to initialize the chunked_parquet_writer_options::builder(sink) and
underlying writer before calling ch_in->receive(); then handle the case where
ch_in closes with no messages by returning success instead of throwing. Locate
symbols write_parquet, ch_in->receive(), streaming::TableChunk, and
chunked_parquet_writer_options::builder(sink) to implement: add the new
parameter, branch to initialize writer from the provided schema, and ensure an
empty msg from receive results in a clean return rather than RAPIDSMPF_EXPECTS
failure.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hpp`:
- Around line 13-15: The header parquet_writer.hpp is not self-contained because
write_parquet(...) uses std::vector<std::string> but the file does not include
<string>; add `#include` <string> to the header so it explicitly declares
std::string and guarantees compilation regardless of include order, ensuring the
declaration of write_parquet (and any other uses of std::string in that header)
compiles correctly.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp`:
- Around line 100-116: The code heap-allocates owner with raw new in the filter
builder (variable owner) and only wraps it late with OwningWrapper, risking
leaks if intermediate allocations throw; replace the raw new by creating a local
std::vector<std::any> (or a std::shared_ptr<std::vector<std::any>>) on the stack
or as a smart pointer, populate it, then pass a managed pointer (e.g.,
shared_ptr.get() with a capture or move a shared_ptr into OwningWrapper) to
streaming::Filter so cleanup happens on exceptions; update both helper sites
(the block creating timestamp_scalar, literal, column_name_reference, operation
and the return creating OwningWrapper) to use RAII rather than raw new.
In `@cpp/libcudf_streaming/benchmarks/utils/misc.hpp`:
- Around line 55-71: parse_integer currently validates the parsed long long
(val) only against min_val/max_val typed as int64_t then casts to T, which can
overflow for smaller integer types; update parse_integer to compute the
effective bounds as the intersection of the provided min_val/max_val and
std::numeric_limits<T>::min()/max(), validate val against that intersection
before assigning to output, and produce the same error messages when
out-of-range; reference the function parse_integer, parameters output, str,
min_val, max_val, and the local val when making this change.
In `@cpp/libcudf_streaming/tests/streaming/test_leaf_actor.cpp`:
- Around line 147-148: The atomic accumulators (e.g., the local std::atomic<int>
variables named result that are passed into consumer(ctx, ch, result)) must be
value-initialized before any concurrent use; change their declaration to
initialize them to 0 (e.g., std::atomic<int> result{0};) wherever you declare
them (both occurrences used with consumer and the other consumer call around
those lines) so fetch_add operates on a defined value.
In `@cpp/libcudf_streaming/tests/test_shuffler.cpp`:
- Around line 327-329: The loop over the futures currently calls f.wait(), which
only blocks and hides exceptions from async workers; replace the wait() call in
the for (auto& f : futures) loop with f.get() and keep the ASSERT_NO_THROW
around that call so any exceptions thrown by the worker tasks (from the futures
vector) are propagated and cause the test to fail as intended.
In `@python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py`:
- Line 391: Input discovery uses args.input.glob("**/*") which yields
directories and non-Parquet files and can break read_parquet; change the
discovery to only include regular files with Parquet extensions. Update the code
that builds paths (the variable assigned with paths=sorted(...)) to either use
args.input.glob("**/*.parquet") (and include ".pq" if needed) or filter the
globbed entries by Path.is_file() and by suffix in {".parquet",".pq"} before
converting to str, so downstream calls to read_parquet only receive valid file
paths.
- Around line 205-211: The code is recreating a new BufferResource via
BufferResource(rmm.mr.get_current_device_resource()) which discards the
caller-provided BufferResource (and its --spill-device limits); instead, stop
reconstructing br and pass the original caller-provided BufferResource into the
Shuffler. Locate the BufferResource creation around the Shuffler construction
(symbol names: BufferResource, br, Shuffler) and remove or replace the
BufferResource(...) instantiation so that the existing br instance with
configured spill limits is used when constructing Shuffler.
In `@python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py`:
- Around line 366-369: The code always calls ray.init(address="auto") even when
args.ray_address is supplied; update the connection branch in the ray.init call
to use the actual provided address. In the block that checks args.ray_address or
os.environ.get("RAY_ADDRESS"), compute an address value like address =
args.ray_address or os.environ.get("RAY_ADDRESS") (falling back to "auto" only
if neither is set) and pass that into ray.init(address=address) so explicit
--ray-address is honored; leave the else branch (ray.init(num_gpus=...,
dashboard_host=...)) unchanged.
- Line 372: The current paths=sorted(map(str, args.input.glob("**/*"))) collects
directories and non-Parquet files which can break the Parquet reader; change the
collection to only regular files with Parquet extensions (e.g., use
args.input.rglob("**/*.parquet") or filter args.input.rglob("*") by p.is_file()
and p.suffix in {".parquet", ".parquet.gzip", ".parquet.gz"}), update the
variable referenced as paths so only valid file paths (strings) are produced
before passing to the Parquet reader.
- Around line 209-213: The function read_and_insert can raise UnboundLocalError
because column_names is only set inside the for-loop when read_batch is called;
if paths is empty the loop never runs and the return on column_names fails. Fix
by initializing column_names before the loop (e.g., column_names = [] or None)
and ensure the function still calls insert_finished() and returns that safe
default when no batches were processed; update references in read_and_insert to
use this initialized value so insert_chunk, insert_finished, and the final
return work even when num_workers > num_input_files.
In `@python/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.py`:
- Around line 76-80: The string column generation uses ["cat", "dog"] *
(self._num_rows // 2) which produces only 2*(self._num_rows//2) entries and
breaks for odd self._num_rows; change the construction used in
plc.Column.from_iterable_of_py so it always yields exactly self._num_rows items
(e.g., generate values with a bounded iterator or list comprehension that
repeats "cat"/"dog" for range(self._num_rows) and slices/truncates as needed) so
the column length matches self._num_rows.
In `@python/cudf_streaming/cudf_streaming/examples/streaming_basic_example.py`:
- Around line 45-49: The Context instance created as ctx may not be shut down on
exception; ensure Context.shutdown() always runs by wrapping the actor logic
that uses ctx (the block after ctx = Context(...)) in a try/finally (or convert
Context to a contextmanager) so that ctx.shutdown() is called in the finally
block; locate references to ctx and Context.shutdown() to update both the main
usage and the similar section around lines 129-150 so cleanup runs on both
success and exception paths.
---
Nitpick comments:
In `@cpp/libcudf_streaming/scripts/ndsh.py`:
- Around line 446-448: The zip call over output.column_names, output.columns,
and expected.columns currently uses strict=False which can silently ignore
length mismatches; change the zip to use strict=True so mismatched column counts
raise immediately (since schema equality is already checked at the earlier
verification around line 412). Update the call referencing output.column_names,
output.columns, and expected.columns in ndsh.py to pass strict=True to zip to
enforce length equality during iteration.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 266e2820-2b9c-4ef1-aa1f-03c34dc014d3
📒 Files selected for processing (62)
cpp/libcudf_streaming/CMakeLists.txtcpp/libcudf_streaming/benchmarks/CMakeLists.txtcpp/libcudf_streaming/benchmarks/bench_pack.cppcpp/libcudf_streaming/benchmarks/bench_partition.cppcpp/libcudf_streaming/benchmarks/bench_shuffle.cppcpp/libcudf_streaming/benchmarks/streaming/CMakeLists.txtcpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cppcpp/libcudf_streaming/benchmarks/streaming/data_generator.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/CMakeLists.txtcpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/join.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q01.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q03.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q04.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q09.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q17.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q18.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q21.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hppcpp/libcudf_streaming/benchmarks/utils/misc.hppcpp/libcudf_streaming/benchmarks/utils/random_data.cucpp/libcudf_streaming/benchmarks/utils/random_data.hppcpp/libcudf_streaming/benchmarks/utils/rmm_utils.hppcpp/libcudf_streaming/examples/CMakeLists.txtcpp/libcudf_streaming/examples/example_shuffle.cppcpp/libcudf_streaming/scripts/ndsh.pycpp/libcudf_streaming/tests/CMakeLists.txtcpp/libcudf_streaming/tests/streaming/test_allgather.cppcpp/libcudf_streaming/tests/streaming/test_leaf_actor.cppcpp/libcudf_streaming/tests/streaming/test_shuffler.cppcpp/libcudf_streaming/tests/test_shuffler.cppcpp/libcudf_streaming/tests/test_shuffler_many_streams.cpppython/cudf_streaming/cudf_streaming/examples/__init__.pypython/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.pypython/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.pypython/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.pypython/cudf_streaming/cudf_streaming/examples/streaming_basic_example.pypython/cudf_streaming/cudf_streaming/tests/test_allgather.pypython/cudf_streaming/cudf_streaming/tests/test_integration_partition.pypython/cudf_streaming/cudf_streaming/tests/test_shuffler.pypython/cudf_streaming/cudf_streaming/tests/test_sparse_alltoall.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_allgather.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_define_actor.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_fanout.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_leaf_actor.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_shuffler.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_sparse_alltoall.py
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp`:
- Around line 352-353: The call to harmonic_mean(elapsed_vec) can crash if
elapsed_vec is empty (e.g., user passed -r 0); update argument parsing or the
reporting code to guarantee num_runs >= 1 or to handle the empty vector before
computing the harmonic mean: validate the parsed num_runs in the ArgumentParser
(ensure num_runs >= 1 and return an error/usage message) or, if you prefer
keeping parsing, check elapsed_vec.empty() before calling harmonic_mean in the
reporting block and skip/short-circuit the mean calculation (or set a safe
default/print a warning) to avoid division by zero; reference symbols: num_runs,
ArgumentParser, elapsed_vec, harmonic_mean.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cpp`:
- Line 325: The single CUDA warm-up call cudaFree(nullptr) must be wrapped with
error checking: replace the bare cudaFree(nullptr) with capturing its return
(cudaError_t err = cudaFree(nullptr)); check if err != cudaSuccess and handle it
(log via process/logger or throw) using cudaGetErrorString(err); optionally
handle/ignore specific initialization errors if intended, but do not leave the
call unchecked. Locate the cudaFree(nullptr) call in bench_read.cpp and update
it to use the error-checked pattern so failures are not silently ignored.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cpp`:
- Line 259: The CUDA warm-up call cudaFree(nullptr) currently lacks error
checking; change it to capture the return value (cudaError_t err =
cudaFree(nullptr)) and handle non-success results by logging the error via
cudaGetErrorString(err) and failing gracefully (e.g., print error and
return/exit) where the call appears (the cudaFree(nullptr) call in q01.cpp);
ensure any existing helper macros or functions for CUDA error checking are used
(or add a short check) so CUDA initialization failures are detected and
reported.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp`:
- Around line 243-246: Remove the duplicated cudaFree(nullptr) and add proper
CUDA error checking for the remaining call: replace the two cudaFree(nullptr)
occurrences around the rapidsmpf::ndsh::FinalizeMPI finalize{} block with a
single cudaFree(nullptr) that captures its return value (cudaError_t) and checks
for cudaSuccess, logging or handling the error if it fails (use the same
logging/error-handling pattern used elsewhere in this file). Ensure the single
checked call is located either before or after FinalizeMPI as intended, and
reference/modify the cudaFree(nullptr) invocation in q04.cpp accordingly.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cpp`:
- Line 316: The cuda warm-up call currently uses cudaFree(nullptr) without error
handling; change it to check and handle the returned CUDA error for robustness
by replacing the bare call with the project's CUDA error-checking pattern (e.g.,
use the CUDA_TRY macro or capture the cudaError_t from cudaFree(nullptr) and
call the existing error/log/exit helper). Update the call site where
cudaFree(nullptr) appears so failures are logged/handled consistently with other
CUDA calls in this file (refer to the cudaFree(nullptr) invocation and use
CUDA_TRY or the file's standard cuda error handling routine).
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q17.sql`:
- Line 14: The correlated subquery in q17.sql currently calls
read_parquet('/datasets/toaugspurger/tpch-rs/scale-10/lineitem/*.parquet')
directly which hardcodes dataset and scale; change that to reference the outer
query's lineitem view (use the view name "lineitem" instead of
read_parquet(...)) so the subquery becomes correlated with the outer alias and
uses the provided input directory at runtime; ensure the aliasing and correlated
column references (the outer lineitem.* fields used inside the subquery) remain
correct after replacing the read_parquet call.
In `@cpp/libcudf_streaming/scripts/ndsh.py`:
- Around line 704-717: The code sets results[query_name] = False when
expected_path doesn't exist but then still calls compare_parquet(result_path,
expected_path, ...), causing errors; modify the block around the expected_path
check (the branch that prints "FAILED: Expected file does not exist") to skip
the comparison for that query—e.g., add an early continue/return after setting
results[query_name] = False and printing the message so compare_parquet is not
invoked for that missing expected_path.
In `@python/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.py`:
- Around line 70-81: The string column generation can produce too few elements
when self._num_rows is odd; update the creation passed to
plc.Column.from_iterable_of_py for the string column so it yields exactly
self._num_rows items (for example, generate the repeating pattern and then
slice/truncate to self._num_rows or use itertools.islice over
itertools.cycle(["cat","dog"]) ), ensuring the list/iterable length matches
self._num_rows before constructing the table with plc.Table and
plc.Column.from_iterable_of_py.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: cf5d64af-526f-4cc3-8e07-c5219c59675d
📒 Files selected for processing (62)
cpp/libcudf_streaming/CMakeLists.txtcpp/libcudf_streaming/benchmarks/CMakeLists.txtcpp/libcudf_streaming/benchmarks/bench_pack.cppcpp/libcudf_streaming/benchmarks/bench_partition.cppcpp/libcudf_streaming/benchmarks/bench_shuffle.cppcpp/libcudf_streaming/benchmarks/streaming/CMakeLists.txtcpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cppcpp/libcudf_streaming/benchmarks/streaming/data_generator.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/CMakeLists.txtcpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/join.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q01.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q03.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q04.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q09.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q17.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q18.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q21.sqlcpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hppcpp/libcudf_streaming/benchmarks/utils/misc.hppcpp/libcudf_streaming/benchmarks/utils/random_data.cucpp/libcudf_streaming/benchmarks/utils/random_data.hppcpp/libcudf_streaming/benchmarks/utils/rmm_utils.hppcpp/libcudf_streaming/examples/CMakeLists.txtcpp/libcudf_streaming/examples/example_shuffle.cppcpp/libcudf_streaming/scripts/ndsh.pycpp/libcudf_streaming/tests/CMakeLists.txtcpp/libcudf_streaming/tests/streaming/test_allgather.cppcpp/libcudf_streaming/tests/streaming/test_leaf_actor.cppcpp/libcudf_streaming/tests/streaming/test_shuffler.cppcpp/libcudf_streaming/tests/test_shuffler.cppcpp/libcudf_streaming/tests/test_shuffler_many_streams.cpppython/cudf_streaming/cudf_streaming/examples/__init__.pypython/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.pypython/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.pypython/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.pypython/cudf_streaming/cudf_streaming/examples/streaming_basic_example.pypython/cudf_streaming/cudf_streaming/tests/test_allgather.pypython/cudf_streaming/cudf_streaming/tests/test_integration_partition.pypython/cudf_streaming/cudf_streaming/tests/test_shuffler.pypython/cudf_streaming/cudf_streaming/tests/test_sparse_alltoall.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_allgather.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_define_actor.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_fanout.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_leaf_actor.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_shuffler.pypython/cudf_streaming/cudf_streaming/tests/test_streaming_sparse_alltoall.py
✅ Files skipped from review due to trivial changes (1)
- python/cudf_streaming/cudf_streaming/examples/init.py
🚧 Files skipped from review as they are similar to previous changes (46)
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q09.sql
- cpp/libcudf_streaming/tests/CMakeLists.txt
- cpp/libcudf_streaming/benchmarks/utils/misc.hpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q04.sql
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.hpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.hpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q03.sql
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q01.sql
- python/cudf_streaming/cudf_streaming/tests/test_streaming_sparse_alltoall.py
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q21.sql
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.cpp
- python/cudf_streaming/cudf_streaming/tests/test_streaming_fanout.py
- cpp/libcudf_streaming/benchmarks/streaming/CMakeLists.txt
- cpp/libcudf_streaming/examples/example_shuffle.cpp
- python/cudf_streaming/cudf_streaming/tests/test_streaming_shuffler.py
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/CMakeLists.txt
- cpp/libcudf_streaming/benchmarks/bench_partition.cpp
- cpp/libcudf_streaming/tests/streaming/test_shuffler.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp
- cpp/libcudf_streaming/CMakeLists.txt
- python/cudf_streaming/cudf_streaming/tests/test_integration_partition.py
- python/cudf_streaming/cudf_streaming/tests/test_sparse_alltoall.py
- python/cudf_streaming/cudf_streaming/tests/test_streaming_allgather.py
- cpp/libcudf_streaming/benchmarks/CMakeLists.txt
- cpp/libcudf_streaming/tests/test_shuffler_many_streams.cpp
- cpp/libcudf_streaming/benchmarks/utils/rmm_utils.hpp
- cpp/libcudf_streaming/benchmarks/utils/random_data.hpp
- python/cudf_streaming/cudf_streaming/tests/test_streaming_define_actor.py
- python/cudf_streaming/cudf_streaming/tests/test_allgather.py
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.hpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp
- cpp/libcudf_streaming/benchmarks/streaming/data_generator.hpp
- cpp/libcudf_streaming/tests/streaming/test_allgather.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.cpp
- cpp/libcudf_streaming/tests/streaming/test_leaf_actor.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp
- cpp/libcudf_streaming/benchmarks/utils/random_data.cu
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.hpp
- python/cudf_streaming/cudf_streaming/examples/streaming_basic_example.py
- cpp/libcudf_streaming/benchmarks/bench_pack.cpp
- python/cudf_streaming/cudf_streaming/tests/test_shuffler.py
- python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py
- cpp/libcudf_streaming/tests/test_shuffler.cpp
- python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py
Update includes and namespace references in tests, benchmarks, and
examples to use cudf_streaming headers and namespaces instead of the
removed rapidsmpf integrations headers:
- rapidsmpf/integrations/cudf/partition.hpp -> cudf_streaming/integrations/partition.hpp
- rapidsmpf/integrations/cudf/bloom_filter.hpp -> cudf_streaming/streaming/bloom_filter.hpp
- rapidsmpf/streaming/cudf/{table_chunk,partition,parquet,bloom_filter}.hpp ->
cudf_streaming/streaming/{table_chunk,partition,parquet,bloom_filter}.hpp
- rapidsmpf::{partition_and_pack,partition_and_split,unpack_and_concat,unspill_partitions} ->
cudf_streaming::integrations::*
- rapidsmpf::streaming::{TableChunk,BloomFilter,Filter,to_message} ->
cudf_streaming::streaming::*
- rapidsmpf::streaming::actor::{partition_and_pack,unpack_and_concat,read_parquet} ->
cudf_streaming::streaming::actor::*
- Add BloomFilter::fitting_num_blocks static method
- Add --expt-extended-lambda for CUDA benchmark compilation
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp (1)
488-490:⚠️ Potential issue | 🔴 CriticalCRITICAL: Unchecked CUDA warm-up call (
cudaFree(nullptr))
cudaFree(nullptr)inmain()is not error-checked; if CUDA init fails, the failure is silently ignored beforecudf::initialize(), unlike the other CUDA calls inmain().Suggested fix
- cudaFree(nullptr); + RAPIDSMPF_CUDA_TRY(cudaFree(nullptr));🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp` around lines 488 - 490, The cudaFree(nullptr) warm-up call in main() is currently unchecked; replace it with an explicit cudaError_t check so failures are detected before calling cudf::initialize(): call cudaFree(nullptr), store the returned cudaError_t, and if it is not cudaSuccess log the error (using cudaGetErrorString) and abort/exit non‑zero (or throw) so initialization does not proceed; update the block around rapidsmpf::ndsh::FinalizeMPI finalize{} / cudaFree(nullptr) / cudf::initialize() to perform this check and early exit on error.Source: Coding guidelines
♻️ Duplicate comments (2)
cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp (1)
243-246:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winCRITICAL: Check the CUDA warm-up call(s).
Both
cudaFree(nullptr)calls ignore CUDA failures during startup. If both placements are intentional for the MPI workaround, they still needRAPIDSMPF_CUDA_TRY; otherwise the extra call should be removed in the same change.Suggested fix
- cudaFree(nullptr); + RAPIDSMPF_CUDA_TRY(cudaFree(nullptr)); rapidsmpf::ndsh::FinalizeMPI finalize{}; - cudaFree(nullptr); + RAPIDSMPF_CUDA_TRY(cudaFree(nullptr));As per coding guidelines, every CUDA call must have error checking.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp` around lines 243 - 246, The two CUDA warm-up calls using cudaFree(nullptr) currently have no error checking; locate the cudaFree(nullptr) invocations (around the rapidsmpf::ndsh::FinalizeMPI finalize{} block) and either remove the redundant call if it was accidentally duplicated or wrap each intentional cudaFree(nullptr) with the RAPIDSMPF_CUDA_TRY macro to ensure failures are caught and reported; make the change so every CUDA call in this area is checked consistently.Source: Coding guidelines
cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp (1)
220-242:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winCRITICAL: Release this one-shot latch exactly once.
Lines 554-559 create
latch(1), but Line 241 decrements it for every grouped chunk. That can fail on the second chunk, and if no chunk arrives the secondread_lineitemstays blocked forever.Suggested fix
rapidsmpf::streaming::Actor filter_grouped_greater( std::shared_ptr<rapidsmpf::streaming::Context> ctx, std::shared_ptr<rapidsmpf::streaming::Channel> ch_in, std::shared_ptr<rapidsmpf::streaming::Channel> ch_out, std::shared_ptr<coro::latch> latch) { rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; + bool latch_released = false; while (!ch_out->is_shutdown()) { auto msg = co_await ch_in->receive(); if (msg.empty()) { break; } auto chunk = co_await msg.release<cudf_streaming::streaming::TableChunk>().make_available(ctx); @@ - latch->count_down(); + if (!latch_released) { + latch->count_down(); + latch_released = true; + } co_await ch_out->send(cudf_streaming::streaming::to_message( msg.sequence_number(), std::make_unique<cudf_streaming::streaming::TableChunk>( cudf::apply_boolean_mask( chunk.table_view().select({0}), mask->view(), chunk.stream(), ctx->br()->device_mr()), chunk.stream()))); } + if (!latch_released) { latch->count_down(); } co_await ch_out->drain(ctx->executor()); }Also applies to: 553-559
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp` around lines 220 - 242, The latch is currently being counted down inside the per-chunk loop in filter_grouped_greater which can decrement it multiple times and leave downstream readers blocked; change the logic so the one-shot latch created (latch(1)) is decremented exactly once — for example, move or guard latch->count_down() so it runs only on the first successful group processing (or immediately after the first co_await ch_out->send of the grouped stream), or defer the single count-down to after the loop exits/when you know no further chunks will be produced; update filter_grouped_greater to perform that single-count decrement (or add a boolean flag like seen_first_group to ensure latch->count_down() is called only once) so the latch semantics match its one-shot use.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp`:
- Around line 488-490: The cudaFree(nullptr) warm-up call in main() is currently
unchecked; replace it with an explicit cudaError_t check so failures are
detected before calling cudf::initialize(): call cudaFree(nullptr), store the
returned cudaError_t, and if it is not cudaSuccess log the error (using
cudaGetErrorString) and abort/exit non‑zero (or throw) so initialization does
not proceed; update the block around rapidsmpf::ndsh::FinalizeMPI finalize{} /
cudaFree(nullptr) / cudf::initialize() to perform this check and early exit on
error.
---
Duplicate comments:
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp`:
- Around line 243-246: The two CUDA warm-up calls using cudaFree(nullptr)
currently have no error checking; locate the cudaFree(nullptr) invocations
(around the rapidsmpf::ndsh::FinalizeMPI finalize{} block) and either remove the
redundant call if it was accidentally duplicated or wrap each intentional
cudaFree(nullptr) with the RAPIDSMPF_CUDA_TRY macro to ensure failures are
caught and reported; make the change so every CUDA call in this area is checked
consistently.
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp`:
- Around line 220-242: The latch is currently being counted down inside the
per-chunk loop in filter_grouped_greater which can decrement it multiple times
and leave downstream readers blocked; change the logic so the one-shot latch
created (latch(1)) is decremented exactly once — for example, move or guard
latch->count_down() so it runs only on the first successful group processing (or
immediately after the first co_await ch_out->send of the grouped stream), or
defer the single count-down to after the loop exits/when you know no further
chunks will be produced; update filter_grouped_greater to perform that
single-count decrement (or add a boolean flag like seen_first_group to ensure
latch->count_down() is called only once) so the latch semantics match its
one-shot use.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d66e9da9-95c5-4f6e-b137-3f0670d88bf5
📒 Files selected for processing (23)
cpp/libcudf_streaming/benchmarks/CMakeLists.txtcpp/libcudf_streaming/benchmarks/bench_partition.cppcpp/libcudf_streaming/benchmarks/bench_shuffle.cppcpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cppcpp/libcudf_streaming/benchmarks/streaming/data_generator.hppcpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hppcpp/libcudf_streaming/examples/example_shuffle.cppcpp/libcudf_streaming/include/cudf_streaming/streaming/bloom_filter.hppcpp/libcudf_streaming/tests/streaming/test_leaf_actor.cppcpp/libcudf_streaming/tests/streaming/test_shuffler.cppcpp/libcudf_streaming/tests/test_shuffler.cpp
🚧 Files skipped from review as they are similar to previous changes (20)
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp
- cpp/libcudf_streaming/benchmarks/CMakeLists.txt
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.cpp
- cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp
- cpp/libcudf_streaming/benchmarks/streaming/data_generator.hpp
- cpp/libcudf_streaming/benchmarks/bench_partition.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cpp
- cpp/libcudf_streaming/tests/streaming/test_leaf_actor.cpp
- cpp/libcudf_streaming/tests/test_shuffler.cpp
- cpp/libcudf_streaming/examples/example_shuffle.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cpp
- cpp/libcudf_streaming/tests/streaming/test_shuffler.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp
- cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp
There was a problem hiding this comment.
♻️ Duplicate comments (1)
cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp (1)
236-253:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winCRITICAL: Remap merge keys after projecting
values.
partialsare built invaluesorder, but Line 253 still passes the originalkeysindices intocudf::merge. That re-sorts the merged rows by the wrong columns, so this pipeline can return an incorrect top-10.As per coding guidelines, logic errors producing wrong results in
cpp/**must be called out.Suggested fix
std::vector<cudf::table_view> views; std::ranges::transform(partials, std::back_inserter(views), [](auto& t) { return t->view(); }); - auto merged = cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr()); + std::vector<cudf::size_type> merge_keys; + merge_keys.reserve(keys.size()); + for (auto const key : keys) { + auto const it = std::find(values.begin(), values.end(), key); + RAPIDSMPF_EXPECTS(it != values.end(), "All sort keys must be projected into `values`"); + merge_keys.push_back(static_cast<cudf::size_type>(std::distance(values.begin(), it))); + } + auto merged = cudf::merge(views, merge_keys, order, {}, out_stream, ctx->br()->device_mr());🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp` around lines 236 - 253, partials currently hold projected `values` in the order produced by per-chunk `indices` (from cudf::sorted_order) but cudf::merge is called with the original `keys`, so rows get re-sorted by the wrong columns; fix by also building and passing the remapped key tables/views that correspond to the same `indices` used to build `partials` (i.e., for each chunk, gather chunk.table_view().select(keys) with the same split(indices->view(), {k}, chunk.stream()).front() and push those key-partials into a separate vector, transform them into views like you do for `partials`, and pass that vector of key views (instead of the original `keys`) to cudf::merge); ensure you use the same streams and device_mr (ctx->br()->device_mr()) when creating these gathered key-partials.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp`:
- Around line 236-253: partials currently hold projected `values` in the order
produced by per-chunk `indices` (from cudf::sorted_order) but cudf::merge is
called with the original `keys`, so rows get re-sorted by the wrong columns; fix
by also building and passing the remapped key tables/views that correspond to
the same `indices` used to build `partials` (i.e., for each chunk, gather
chunk.table_view().select(keys) with the same split(indices->view(), {k},
chunk.stream()).front() and push those key-partials into a separate vector,
transform them into views like you do for `partials`, and pass that vector of
key views (instead of the original `keys`) to cudf::merge); ensure you use the
same streams and device_mr (ctx->br()->device_mr()) when creating these gathered
key-partials.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3c75adf6-e661-43df-8540-b35fcf31f103
📒 Files selected for processing (9)
cpp/libcudf_streaming/benchmarks/bench_partition.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cppcpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hppcpp/src/io/utilities/config_utils.cpp
🚧 Files skipped from review as they are similar to previous changes (7)
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cpp
- cpp/libcudf_streaming/benchmarks/bench_partition.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp
- cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp
mroeschke
approved these changes
Jun 9, 2026
bdice
requested changes
Jun 9, 2026
bdice
left a comment
bdice
left a comment
Contributor
There was a problem hiding this comment.
One change needed in devcontainers.
| sccache --zero-stats; | ||
| build-all -j0 -DBUILD_BENCHMARKS=ON --verbose 2>&1 | tee telemetry-artifacts/build.log; | ||
| configure-cudf-cpp -DBUILD_BENCHMARKS=ON; | ||
| configure-cudf_streaming-cpp -DBUILD_BENCHMARKS=OFF; |
Contributor
There was a problem hiding this comment.
We enable -DBUILD_BENCHMARKS in rapidsmpf. https://github.com/rapidsai/rapidsmpf/blob/5234785a0255847784770948b9728beafe00a8a7/.github/workflows/pr.yaml#L356
The fix is probably switching to rapidsai/devcontainers:26.08-cpp-cuda13.2-ucx1.19.0-openmpi5.0.10, let's do that (and similar for 12.9). https://github.com/rapidsai/rapidsmpf/blob/5234785a0255847784770948b9728beafe00a8a7/.devcontainer/cuda13.2-pip/devcontainer.json#L8
Then delete these special configure-* commands.
Contributor
There was a problem hiding this comment.
Update: this isn't viable. We switched to using rapidsmpf wheels, which don't have MPI support. Therefore, cudf-streaming benchmarks can't be built for pip devcontainers.
Follow-ups:
- Enable all benchmarks for conda devcontainers
- Split up (CMake?) logic so that MPI benchmarks and UCX benchmarks can be built separately
- Enable UCX benchmarks for pip devcontainers (requires split above)
nirandaperera
left a comment
nirandaperera
left a comment
Contributor
There was a problem hiding this comment.
Seems like we can remove some of the tests in the files. Maybe we can do that later.
| @@ -0,0 +1,71 @@ | |||
| /** | |||
Contributor
There was a problem hiding this comment.
Is there a better way to share these headers in the benchmarks/utils dir? I think some of these utils needs to/ could be shared by both rapidsmpf and cudf streaming.
Maybe move to rapidsmpf/utils dir?
Comment on lines
+50
to
+73
| std::unique_ptr<cudf::column> random_column(cudf::size_type nrows, | ||
| std::int32_t min_val, | ||
| std::int32_t max_val, | ||
| rmm::cuda_stream_view stream, | ||
| rmm::device_async_resource_ref mr) | ||
| { | ||
| auto vec = | ||
| random_device_vector(rapidsmpf::safe_cast<std::size_t>(nrows), min_val, max_val, stream, mr); | ||
| return std::make_unique<cudf::column>(std::move(vec), rmm::device_buffer{0, stream, mr}, 0); | ||
| } | ||
|
|
||
| cudf::table random_table(cudf::size_type ncolumns, | ||
| cudf::size_type nrows, | ||
| std::int32_t min_val, | ||
| std::int32_t max_val, | ||
| rmm::cuda_stream_view stream, | ||
| rmm::device_async_resource_ref mr) | ||
| { | ||
| std::vector<std::unique_ptr<cudf::column>> cols; | ||
| for (auto i = 0; i < ncolumns; ++i) { | ||
| cols.push_back(random_column(nrows, min_val, max_val, stream, mr)); | ||
| } | ||
| return cudf::table(std::move(cols)); | ||
| } |
Contributor
There was a problem hiding this comment.
These are the only utils related to cudf
Comment on lines
+164
to
+223
| class ShufflerAsyncTest | ||
| : public BaseStreamingShuffle, | ||
| public ::testing::WithParamInterface<std::tuple<std::size_t, std::uint32_t>> { | ||
| protected: | ||
| std::size_t n_inserts; | ||
| std::uint32_t n_partitions; | ||
|
|
||
| static constexpr OpID op_id = 0; | ||
| static constexpr std::size_t n_elements = 100; | ||
|
|
||
| void SetUp() override | ||
| { | ||
| std::tie(n_inserts, n_partitions) = GetParam(); | ||
|
|
||
| BaseStreamingShuffle::SetUpWithThreads(4); | ||
| } | ||
|
|
||
| void TearDown() override { BaseStreamingShuffle::TearDown(); } | ||
| }; | ||
|
|
||
| INSTANTIATE_TEST_SUITE_P(StreamingShuffler, | ||
| ShufflerAsyncTest, | ||
| ::testing::Combine(::testing::Values(1, 10), // number of inserts | ||
| ::testing::Values(1, 10, 100) // number of partitions | ||
| ), | ||
| [](const testing::TestParamInfo<ShufflerAsyncTest::ParamType>& info) { | ||
| return "ninserts_" + std::to_string(std::get<0>(info.param)) + | ||
| "_nparts_" + std::to_string(std::get<1>(info.param)); | ||
| }); | ||
|
|
||
| TEST_P(ShufflerAsyncTest, insert_wait_extract) | ||
| { | ||
| auto comm = GlobalEnvironment->comm_; | ||
| auto shuffler = std::make_unique<ShufflerAsync>(ctx, comm, op_id, n_partitions); | ||
|
|
||
| for (std::size_t i = 0; i < n_inserts; ++i) { | ||
| std::unordered_map<shuffler::PartID, PackedData> data; | ||
| data.reserve(n_partitions); | ||
| for (shuffler::PartID pid = 0; pid < n_partitions; ++pid) { | ||
| data.emplace(pid, generate_packed_data(n_elements, 0, stream, *br)); | ||
| } | ||
| shuffler->insert(std::move(data)); | ||
| } | ||
|
|
||
| coro::sync_wait(shuffler->insert_finished()); | ||
|
|
||
| auto local_pids = | ||
| shuffler::Shuffler::local_partitions(comm, n_partitions, &shuffler::Shuffler::round_robin); | ||
|
|
||
| std::vector<shuffler::PartID> finished_pids; | ||
| std::size_t n_chunks_received = 0; | ||
| for (auto pid : local_pids) { | ||
| auto chunks = shuffler->extract(pid); | ||
| n_chunks_received += chunks.size(); | ||
| finished_pids.push_back(pid); | ||
| } | ||
|
|
||
| EXPECT_EQ(n_inserts * local_pids.size() * comm->nranks(), n_chunks_received); | ||
| EXPECT_EQ(local_pids, finished_pids); | ||
| } |
Comment on lines
+34
to
+56
| TEST_F(BaseStreamingShuffle, zero_owned_partitions_completes) | ||
| { | ||
| auto comm = GlobalEnvironment->comm_; | ||
| if (comm->nranks() < 2) { | ||
| GTEST_SKIP() << "Need at least 2 ranks so that some rank owns 0 partitions"; | ||
| } | ||
| constexpr Rank owner = 0; | ||
| auto collapse = [](std::shared_ptr<Communicator> const&, | ||
| shuffler::PartID, | ||
| shuffler::PartID) -> Rank { return owner; }; | ||
| constexpr OpID op_id = 0; | ||
| constexpr shuffler::PartID total = 4; | ||
| auto shuffler = std::make_unique<ShufflerAsync>(ctx, comm, op_id, total, collapse); | ||
|
|
||
| coro::sync_wait(shuffler->insert_finished()); | ||
|
|
||
| auto local_pids = shuffler->local_partitions(); | ||
| if (comm->rank() == owner) { | ||
| EXPECT_EQ(local_pids.size(), total); | ||
| } else { | ||
| EXPECT_TRUE(local_pids.empty()); | ||
| } | ||
| } |
| @@ -0,0 +1,171 @@ | |||
| /** | |||
Comment on lines
+33
to
+103
| TEST(ReceivedChunks, spill_skips_control_messages) | ||
| { | ||
| auto mr = cudf::get_current_device_resource_ref(); | ||
| auto br = rapidsmpf::BufferResource::create(mr); | ||
|
|
||
| rapidsmpf::shuffler::detail::ReceivedChunks received; | ||
|
|
||
| // Control messages have no data buffer (data_ == nullptr); spill must skip them | ||
| // rather than calling data_memory_type(), which throws if data_ is null. | ||
| received.insert(rapidsmpf::shuffler::detail::Chunk::from_finished_partition( | ||
| /*chunk_id=*/0, /*part_id=*/0, /*expected_num_chunks=*/1)); | ||
|
|
||
| EXPECT_EQ(received.spill(br.get(), /*amount=*/1024), 0UL); | ||
| } | ||
|
|
||
| TEST(ReceivedChunks, spill_respects_amount) | ||
| { | ||
| auto mr = cudf::get_current_device_resource_ref(); | ||
| auto br = rapidsmpf::BufferResource::create(mr); | ||
| auto stream = cudf::get_default_stream(); | ||
|
|
||
| rapidsmpf::shuffler::detail::ReceivedChunks received; | ||
| constexpr std::size_t chunk_size = 100; | ||
|
|
||
| for (rapidsmpf::shuffler::PartID pid = 0; pid < 2; ++pid) { | ||
| auto metadata = std::make_unique<std::vector<std::uint8_t>>(std::size_t{1}, std::uint8_t{0}); | ||
| auto res = br->reserve_or_fail(chunk_size, rapidsmpf::MemoryType::DEVICE); | ||
| auto data = br->make_buffer(chunk_size, stream, res); | ||
| received.insert(rapidsmpf::shuffler::detail::Chunk::from_packed_data( | ||
| 0, pid, rapidsmpf::PackedData{std::move(metadata), std::move(data)})); | ||
| } | ||
|
|
||
| // Two partitions, one 100-byte chunk each. spill() must stop after the first | ||
| // partition satisfies the request; the outer loop must not continue into partition 1. | ||
| EXPECT_EQ(received.spill(br.get(), chunk_size), chunk_size); | ||
| } | ||
|
|
||
| TEST(MetadataMessage, round_trip) | ||
| { | ||
| auto stream = cudf::get_default_stream(); | ||
| auto mr = cudf::get_current_device_resource_ref(); | ||
| auto br = rapidsmpf::BufferResource::create(mr); | ||
|
|
||
| auto metadata = iota_vector<std::uint8_t>(100); | ||
|
|
||
| auto expect = rapidsmpf::shuffler::detail::Chunk::from_packed_data( | ||
| 1, // chunk_id | ||
| 2, // part_id | ||
| rapidsmpf::PackedData{ | ||
| std::make_unique<std::vector<std::uint8_t>>(metadata), // non-empty metadata | ||
| br->move(std::make_unique<rmm::device_buffer>(), stream) // empty data | ||
| }); | ||
|
|
||
| // Extract the metadata from then chunk. | ||
| auto msg = expect.serialize(); | ||
| EXPECT_FALSE(expect.is_metadata_buffer_set()); | ||
|
|
||
| // Create a new chunk by deserializing the message. | ||
| auto result = rapidsmpf::shuffler::detail::Chunk::deserialize(*msg, br.get()); | ||
|
|
||
| EXPECT_TRUE(expect.data_size() == 0 || result.is_data_buffer_set()); | ||
| // They should be identical. | ||
| EXPECT_EQ(expect.part_id(), result.part_id()); | ||
| EXPECT_EQ(expect.chunk_id(), result.chunk_id()); | ||
| EXPECT_EQ(expect.expected_num_chunks(), result.expected_num_chunks()); | ||
| EXPECT_EQ(expect.data_size(), result.data_size()); | ||
| EXPECT_EQ(expect.metadata_size(), result.metadata_size()); | ||
|
|
||
| // The metadata should be identical to the original. | ||
| EXPECT_EQ(metadata, *result.release_metadata_buffer()); | ||
| } |
Comment on lines
+458
to
+643
| TEST(FinishCounterTests, zero_local_partitions_immediately_finished) | ||
| { | ||
| rapidsmpf::shuffler::detail::FinishCounter finish_counter( | ||
| /*nranks=*/2, /*n_local_partitions=*/0); | ||
|
|
||
| EXPECT_TRUE(finish_counter.all_finished()); | ||
| } | ||
|
|
||
| TEST(FinishCounterTests, nonzero_local_partitions_finishes_after_all_chunks) | ||
| { | ||
| rapidsmpf::shuffler::detail::FinishCounter finish_counter( | ||
| /*nranks=*/1, /*n_local_partitions=*/2); | ||
|
|
||
| EXPECT_FALSE(finish_counter.all_finished()); | ||
|
|
||
| // One rank sends 3 chunks total. | ||
| finish_counter.move_goalpost(0, 3); | ||
| finish_counter.add_finished_chunk(); | ||
| finish_counter.add_finished_chunk(); | ||
| EXPECT_FALSE(finish_counter.all_finished()); | ||
|
|
||
| finish_counter.add_finished_chunk(); | ||
| EXPECT_TRUE(finish_counter.all_finished()); | ||
| } | ||
|
|
||
| TEST(FinishCounterTests, multi_rank_completion) | ||
| { | ||
| auto comm = GlobalEnvironment->comm_; | ||
|
|
||
| if (comm->rank() != 0) { GTEST_SKIP() << "Test only runs on rank 0"; } | ||
|
|
||
| // Use nranks partitions so each rank owns exactly 1 partition (round robin). | ||
| auto out_nparts = rapidsmpf::safe_cast<rapidsmpf::shuffler::PartID>(comm->nranks()); | ||
|
|
||
| auto local_partitions = rapidsmpf::shuffler::Shuffler::local_partitions( | ||
| comm, out_nparts, &rapidsmpf::shuffler::Shuffler::round_robin); | ||
| ASSERT_EQ(local_partitions.size(), 1); | ||
|
|
||
| rapidsmpf::shuffler::detail::FinishCounter finish_counter(comm->nranks(), | ||
| local_partitions.size()); | ||
|
|
||
| // Not finished yet. | ||
| EXPECT_FALSE(finish_counter.all_finished()); | ||
|
|
||
| // For nranks ranks, each rank sends 1 data chunk + 1 control, so | ||
| // move_goalpost(rank, 2) per rank. | ||
| for (rapidsmpf::Rank r = 0; r < comm->nranks(); r++) { | ||
| finish_counter.move_goalpost(r, 2); | ||
| } | ||
|
|
||
| // Add finished chunks: 1 data chunk per rank + 1 control per rank = 2 * nranks | ||
| for (rapidsmpf::Rank r = 0; r < comm->nranks(); r++) { | ||
| finish_counter.add_finished_chunk(); // data chunk | ||
| finish_counter.add_finished_chunk(); // control chunk | ||
| } | ||
|
|
||
| EXPECT_TRUE(finish_counter.all_finished()); | ||
| } | ||
|
|
||
| class FinishCounterMultithreadingTest | ||
| : public ::testing::TestWithParam<std::tuple<rapidsmpf::shuffler::PartID, std::uint32_t>> { | ||
| protected: | ||
| rapidsmpf::Rank const nranks{1}; // simulate a single rank | ||
|
|
||
| std::unique_ptr<rapidsmpf::shuffler::detail::FinishCounter> finish_counter; | ||
| rapidsmpf::shuffler::PartID npartitions; | ||
| std::uint32_t nthreads; | ||
|
|
||
| void SetUp() override | ||
| { | ||
| std::tie(npartitions, nthreads) = GetParam(); | ||
|
|
||
| finish_counter = | ||
| std::make_unique<rapidsmpf::shuffler::detail::FinishCounter>(nranks, npartitions); | ||
| } | ||
|
|
||
| void produce_data() | ||
| { | ||
| // Simulate nranks=1: one rank reports chunk count = npartitions + 1 | ||
| // (one data chunk per partition + 1 control message) | ||
| finish_counter->move_goalpost(rapidsmpf::Rank{0}, npartitions + 1); | ||
| for (rapidsmpf::shuffler::PartID i = 0; i <= npartitions; i++) { | ||
| finish_counter->add_finished_chunk(); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| // Parametrize on number of partitions and number of consumer threads | ||
| INSTANTIATE_TEST_SUITE_P(FinishCounterMultithreadingTestP, | ||
| FinishCounterMultithreadingTest, | ||
| testing::Combine(testing::Values(1, 2, 100, 101), | ||
| testing::Values(1, 2, 3)), | ||
| [](const auto& info) { | ||
| return "npartitions_" + std::to_string(std::get<0>(info.param)) + | ||
| "__nthreads_" + std::to_string(std::get<1>(info.param)); | ||
| }); | ||
|
|
||
| TEST_P(FinishCounterMultithreadingTest, concurrent_all_finished_check) | ||
| { | ||
| produce_data(); | ||
|
|
||
| std::atomic<std::uint32_t> n_checks{0}; | ||
| std::vector<std::future<void>> futures; | ||
| for (std::uint32_t tid = 0; tid < nthreads; tid++) { | ||
| futures.emplace_back(std::async(std::launch::async, [&, tid] { | ||
| for (std::uint32_t i = tid; i < npartitions; i += nthreads) { | ||
| EXPECT_TRUE(finish_counter->all_finished()); | ||
| n_checks.fetch_add(1, std::memory_order_relaxed); | ||
| } | ||
| })); | ||
| } | ||
|
|
||
| EXPECT_NO_THROW(std::ranges::for_each(futures, [](auto& f) { f.get(); })); | ||
|
|
||
| EXPECT_EQ(npartitions, n_checks); | ||
| EXPECT_TRUE(finish_counter->all_finished()); | ||
| } | ||
|
|
||
| class ContiguousPartitionAssignmentTest | ||
| : public ::testing::TestWithParam<rapidsmpf::shuffler::PartID> { | ||
| protected: | ||
| void SetUp() override | ||
| { | ||
| comm = GlobalEnvironment->comm_; | ||
| nranks = comm->nranks(); | ||
| rank = comm->rank(); | ||
| total_num_partitions = GetParam(); | ||
| } | ||
|
|
||
| std::shared_ptr<rapidsmpf::Communicator> comm; | ||
| rapidsmpf::Rank nranks; | ||
| rapidsmpf::Rank rank; | ||
| rapidsmpf::shuffler::PartID total_num_partitions; | ||
| }; | ||
|
|
||
| INSTANTIATE_TEST_SUITE_P(PartitionAssignment, | ||
| ContiguousPartitionAssignmentTest, | ||
| testing::Values(1, 2, 3, 5, 7, 10, 16, 100), | ||
| [](const testing::TestParamInfo<rapidsmpf::shuffler::PartID>& info) { | ||
| return "nparts_" + std::to_string(info.param); | ||
| }); | ||
|
|
||
| TEST_P(ContiguousPartitionAssignmentTest, contiguous) | ||
| { | ||
| std::vector<std::vector<rapidsmpf::shuffler::PartID>> rank_partitions(nranks); | ||
| for (rapidsmpf::shuffler::PartID pid = 0; pid < total_num_partitions; ++pid) { | ||
| auto owner = rapidsmpf::shuffler::Shuffler::contiguous(comm, pid, total_num_partitions); | ||
| EXPECT_GE(owner, 0); | ||
| EXPECT_LT(owner, nranks); | ||
| rank_partitions[owner].push_back(pid); | ||
| } | ||
|
|
||
| // Each rank's partitions must be contiguous. | ||
| for (rapidsmpf::Rank r = 0; r < nranks; ++r) { | ||
| auto const& pids = rank_partitions[r]; | ||
| for (std::size_t i = 1; i < pids.size(); ++i) { | ||
| EXPECT_EQ(pids[i], pids[i - 1] + 1); | ||
| } | ||
| } | ||
|
|
||
| // Concatenating all rank partitions should cover [0, total_num_partitions). | ||
| std::vector<rapidsmpf::shuffler::PartID> all_pids; | ||
| for (auto const& pids : rank_partitions) { | ||
| all_pids.insert(all_pids.end(), pids.begin(), pids.end()); | ||
| } | ||
| EXPECT_EQ(all_pids, iota_vector<rapidsmpf::shuffler::PartID>(total_num_partitions)); | ||
| } | ||
|
|
||
| TEST(Shuffler, ShutdownWhilePaused) | ||
| { | ||
| auto progress_thread = GlobalEnvironment->comm_->progress_thread(); | ||
| auto mr = cudf::get_current_device_resource_ref(); | ||
|
|
||
| auto br = rapidsmpf::BufferResource::create(mr); | ||
|
|
||
| auto shuffler = rapidsmpf::shuffler::Shuffler(GlobalEnvironment->comm_, 0, 1, br.get()); | ||
|
|
||
| progress_thread->pause(); | ||
| EXPECT_FALSE(progress_thread->is_running()); | ||
| shuffler.insert_finished(); | ||
| // Progress thread must be running before shuffle shutdown, otherwise we have some | ||
| // orphan messages in the shuffle that are never sent/received. | ||
| progress_thread->resume(); | ||
| EXPECT_TRUE(progress_thread->is_running()); | ||
| EXPECT_NO_THROW(shuffler.shutdown()); | ||
| } |
Comment on lines
+659
to
+778
| class ExtractEmptyPartitionsTest : public cudf::test::BaseFixture { | ||
| public: | ||
| static constexpr rapidsmpf::shuffler::PartID nparts = 10; | ||
| static constexpr auto wait_timeout = std::chrono::seconds(30); | ||
|
|
||
| void SetUp() override | ||
| { | ||
| stream = cudf::get_default_stream(); | ||
| br = rapidsmpf::BufferResource::create(mr()); | ||
|
|
||
| shuffler = std::make_unique<rapidsmpf::shuffler::Shuffler>( | ||
| GlobalEnvironment->comm_, 0, nparts, br.get()); | ||
| } | ||
|
|
||
| void TearDown() override { shuffler.reset(); } | ||
|
|
||
| void insert_chunks( | ||
| std::unordered_map<rapidsmpf::shuffler::PartID, rapidsmpf::PackedData>&& chunks) | ||
| { | ||
| if (!chunks.empty()) { shuffler->insert(std::move(chunks)); } | ||
| shuffler->insert_finished(); | ||
| } | ||
|
|
||
| void verify_extracted_chunks(auto expected_empty_fn) | ||
| { | ||
| EXPECT_NO_THROW(shuffler->wait(wait_timeout)); | ||
| for (auto pid : shuffler->local_partitions()) { | ||
| SCOPED_TRACE("pid: " + std::to_string(pid)); | ||
| std::vector<rapidsmpf::PackedData> chunks; | ||
| EXPECT_NO_THROW({ chunks = shuffler->extract(pid); }); | ||
|
|
||
| if (expected_empty_fn(pid)) { | ||
| EXPECT_TRUE(chunks.empty()); | ||
| } else { | ||
| EXPECT_EQ(GlobalEnvironment->comm_->nranks(), chunks.size()); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| auto empty_packed_data() | ||
| { | ||
| return rapidsmpf::PackedData{std::make_unique<std::vector<std::uint8_t>>(), | ||
| br->move(std::make_unique<rmm::device_buffer>(), stream)}; | ||
| } | ||
|
|
||
| auto non_empty_packed_data() | ||
| { | ||
| return rapidsmpf::PackedData{ | ||
| std::make_unique<std::vector<std::uint8_t>>(10), | ||
| br->move(std::make_unique<rmm::device_buffer>(10, stream), stream)}; | ||
| } | ||
|
|
||
| rmm::cuda_stream_view stream; | ||
| std::shared_ptr<rapidsmpf::BufferResource> br; | ||
| std::unique_ptr<rapidsmpf::shuffler::Shuffler> shuffler; | ||
| }; | ||
|
|
||
| TEST_F(ExtractEmptyPartitionsTest, NoInsertions) | ||
| { | ||
| insert_chunks({}); | ||
| EXPECT_NO_FATAL_FAILURE(verify_extracted_chunks([](auto) { return true; })); | ||
| } | ||
|
|
||
| TEST_F(ExtractEmptyPartitionsTest, AllEmptyInsertions) | ||
| { | ||
| std::unordered_map<rapidsmpf::shuffler::PartID, rapidsmpf::PackedData> chunks; | ||
| for (rapidsmpf::shuffler::PartID pid = 0; pid < nparts; ++pid) { | ||
| chunks.emplace(pid, empty_packed_data()); | ||
| } | ||
|
|
||
| insert_chunks(std::move(chunks)); | ||
| EXPECT_NO_FATAL_FAILURE(verify_extracted_chunks([](auto) { return true; })); | ||
| } | ||
|
|
||
| TEST_F(ExtractEmptyPartitionsTest, SomeEmptyInsertions) | ||
| { | ||
| std::unordered_map<rapidsmpf::shuffler::PartID, rapidsmpf::PackedData> chunks; | ||
| for (rapidsmpf::shuffler::PartID pid = 0; pid < nparts; ++pid) { | ||
| if (pid % 3 == 0) { chunks.emplace(pid, empty_packed_data()); } | ||
| } | ||
|
|
||
| insert_chunks(std::move(chunks)); | ||
| EXPECT_NO_FATAL_FAILURE(verify_extracted_chunks([](auto) { return true; })); | ||
| } | ||
|
|
||
| TEST_F(ExtractEmptyPartitionsTest, SomeEmptyAndNonEmptyInsertions) | ||
| { | ||
| std::unordered_map<rapidsmpf::shuffler::PartID, rapidsmpf::PackedData> chunks; | ||
| for (rapidsmpf::shuffler::PartID pid = 0; pid < nparts; ++pid) { | ||
| if (pid % 3 == 0) { | ||
| chunks.emplace(pid, empty_packed_data()); | ||
| } else { | ||
| chunks.emplace(pid, non_empty_packed_data()); | ||
| } | ||
| } | ||
|
|
||
| insert_chunks(std::move(chunks)); | ||
| EXPECT_NO_FATAL_FAILURE(verify_extracted_chunks([](auto pid) { return pid % 3 == 0; })); | ||
| } | ||
|
|
||
| TEST(ShufflerTest, multiple_shutdowns) | ||
| { | ||
| auto& comm = GlobalEnvironment->comm_; | ||
| auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); | ||
| auto shuffler = | ||
| std::make_unique<rapidsmpf::shuffler::Shuffler>(comm, 0, comm->nranks(), br.get()); | ||
|
|
||
| shuffler->insert_finished(); | ||
| EXPECT_NO_THROW(shuffler->wait(std::chrono::seconds(30))); | ||
| for (auto pid : shuffler->local_partitions()) { | ||
| std::ignore = shuffler->extract(pid); | ||
| } | ||
|
|
||
| constexpr int n_threads = 10; | ||
| std::vector<std::future<void>> futures; | ||
| for (int i = 0; i < n_threads; ++i) { | ||
| futures.emplace_back(std::async(std::launch::async, [&] { shuffler->shutdown(); })); | ||
| } | ||
| std::ranges::for_each(futures, [](auto& future) { future.get(); }); | ||
| } |
bdice
approved these changes
Jun 9, 2026
10 tasks
Contributor
Author
|
/merge |
This was referenced Jun 10, 2026
rapids-bot Bot
pushed a commit
that referenced
this pull request
Jun 10, 2026
#22810 was merged with changes to remove the (un)spill_partition functions from cudf-streaming in favor of using them from RapidsMPF. However, simultaneously #22814 got merged introducing some new use cases for those functions, that allowed #22810 to pass CI and be merged. This change fixes coverage. Authors: - Peter Andreas Entschev (https://github.com/pentschev) - Niranda Perera (https://github.com/nirandaperera) Approvers: - Mads R. B. Kristensen (https://github.com/madsbk) - Niranda Perera (https://github.com/nirandaperera) - Bradley Dice (https://github.com/bdice) URL: #22837
abigalekim
pushed a commit
to abigalekim/cudf
that referenced
this pull request
Jun 12, 2026
NVIDIA#22810 was merged with changes to remove the (un)spill_partition functions from cudf-streaming in favor of using them from RapidsMPF. However, simultaneously NVIDIA#22814 got merged introducing some new use cases for those functions, that allowed NVIDIA#22810 to pass CI and be merged. This change fixes coverage. Authors: - Peter Andreas Entschev (https://github.com/pentschev) - Niranda Perera (https://github.com/nirandaperera) Approvers: - Mads R. B. Kristensen (https://github.com/madsbk) - Niranda Perera (https://github.com/nirandaperera) - Bradley Dice (https://github.com/bdice) URL: NVIDIA#22837
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR adds the tests and benchmarks from rapidsmpf that use cudf types.
Checklist