Skip to content

Commit 2b28962

Browse files
authored
Wait for all submitted host tasks before propagating exceptions in cuIO (#24791)
Closes #24790 Host worker pool futures do not wait for their task on destruction, so when a `get()` throws in cuIO, the remaining tasks can keep running on buffers that unwinding has already freed. This PR adds `cudf::detail::futures_wait_guard`, an RAII guard that waits on all valid futures on scope exit, and uses it across the cuIO readers, writers and host compression paths so submitted tasks never outlive the data they reference. Authors: - Morris (https://github.com/morristai) - Muhammad Haseeb (https://github.com/mhaseeb123) Approvers: - Vukasin Milovanovic (https://github.com/vuule) - Muhammad Haseeb (https://github.com/mhaseeb123) URL: #24791
1 parent 791c138 commit 2b28962

12 files changed

Lines changed: 131 additions & 71 deletions

File tree

‎cpp/include/cudf/detail/utilities/host_worker_pool.hpp‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
#include <BS_thread_pool.hpp>
1313

1414
#include <cstddef>
15+
#include <future>
1516
#include <memory>
1617

1718
namespace cudf::detail {
@@ -106,6 +107,41 @@ class CUDF_EXPORT hierarchical_thread_pool {
106107
[[nodiscard]] std::size_t get_thread_count() const { return pool_.get_thread_count(); }
107108
};
108109

110+
/**
111+
* @brief RAII guard that waits on all valid futures in `futures` upon destruction
112+
*
113+
* Futures returned by `hierarchical_thread_pool::submit_task` do not block on destruction. This
114+
* guard waits for all of them without retrieving their results, so exceptions stored by the tasks
115+
* are discarded and any in-flight exception propagates unchanged.
116+
*
117+
* @note Construct the guard after the futures container and after every object the tasks reference.
118+
*
119+
* @tparam Container Container of `std::future` objects
120+
*/
121+
template <typename Container>
122+
class futures_wait_guard {
123+
public:
124+
/**
125+
* @brief Construct a guard over a container of futures.
126+
*
127+
* @param futures Futures to wait on when the guard is destroyed
128+
*/
129+
explicit futures_wait_guard(Container& futures) noexcept : futures_{futures} {}
130+
131+
futures_wait_guard(futures_wait_guard const&) = delete;
132+
futures_wait_guard& operator=(futures_wait_guard const&) = delete;
133+
134+
~futures_wait_guard() noexcept
135+
{
136+
for (auto& future : futures_) {
137+
if (future.valid()) { future.wait(); }
138+
}
139+
}
140+
141+
private:
142+
Container& futures_;
143+
};
144+
109145
/**
110146
* @brief Retrieves the appropriate thread pool based on the calling thread's context.
111147
*

‎cpp/src/io/comp/compression.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,7 @@ void host_compress(compression_type compression,
332332
auto h_results = cudf::detail::make_pinned_vector<codec_exec_result>(results, stream);
333333

334334
std::vector<std::future<std::pair<size_t, size_t>>> tasks;
335+
cudf::detail::futures_wait_guard tasks_guard{tasks};
335336
auto const num_streams =
336337
std::min<std::size_t>(num_chunks, cudf::detail::host_worker_pool().get_thread_count());
337338
auto const streams = cudf::detail::fork_streams(stream, num_streams);

‎cpp/src/io/comp/decompression.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -571,6 +571,7 @@ void host_decompress(compression_type compression,
571571
stream.sync();
572572

573573
std::vector<std::future<size_t>> tasks;
574+
cudf::detail::futures_wait_guard tasks_guard{tasks};
574575
auto const num_streams =
575576
std::min<std::size_t>(num_chunks, cudf::detail::host_worker_pool().get_thread_count());
576577
auto const streams = cudf::detail::fork_streams(stream, num_streams);

‎cpp/src/io/json/read_json.cu‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
#include <cudf/detail/utilities/buffer_factories.hpp>
1313
#include <cudf/detail/utilities/cuda_memcpy.hpp>
1414
#include <cudf/detail/utilities/getenv_or.hpp>
15+
#include <cudf/detail/utilities/host_worker_pool.hpp>
1516
#include <cudf/detail/utilities/integer_utils.hpp>
1617
#include <cudf/detail/utilities/stream_pool.hpp>
1718
#include <cudf/detail/utilities/vector_factories.hpp>
@@ -754,6 +755,7 @@ device_span<char> ingest_raw_input(device_span<char> buffer,
754755
// delimiter.
755756
auto constexpr num_delimiter_chars = 1;
756757
std::vector<std::future<std::size_t>> thread_tasks;
758+
cudf::detail::futures_wait_guard thread_tasks_guard{thread_tasks};
757759

758760
auto delimiter_map = cudf::detail::make_empty_host_vector<std::size_t>(sources.size(), stream);
759761
std::vector<std::size_t> prefsum_source_sizes(sources.size());
@@ -861,6 +863,7 @@ table_with_metadata read_json_dispatch(host_span<std::unique_ptr<datasource>> so
861863

862864
std::vector<std::unique_ptr<datasource>> compressed_sources;
863865
std::vector<std::future<std::unique_ptr<compressed_host_buffer_source>>> thread_tasks;
866+
cudf::detail::futures_wait_guard thread_tasks_guard{thread_tasks};
864867
auto& thread_pool = pools::tpool();
865868
for (auto& src : sources) {
866869
thread_tasks.emplace_back(thread_pool.submit_task([&reader_opts, &src] {

‎cpp/src/io/orc/writer_impl.cu‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include <cudf/detail/utilities/buffer_factories.hpp>
2222
#include <cudf/detail/utilities/cuda.cuh>
2323
#include <cudf/detail/utilities/cuda_memcpy.hpp>
24+
#include <cudf/detail/utilities/host_worker_pool.hpp>
2425
#include <cudf/detail/utilities/integer_utils.hpp>
2526
#include <cudf/detail/utilities/stream_pool.hpp>
2627
#include <cudf/detail/utilities/vector_factories.hpp>
@@ -2893,6 +2894,7 @@ void writer::impl::write_orc_data_to_sink(encoded_data const& enc_data,
28932894

28942895
// Write stripes
28952896
std::vector<std::future<void>> write_tasks;
2897+
cudf::detail::futures_wait_guard write_tasks_guard{write_tasks};
28962898
for (size_t stripe_id = 0; stripe_id < stripes.size(); ++stripe_id) {
28972899
auto& stripe = stripes[stripe_id];
28982900

‎cpp/src/io/parquet/compact_protocol_reader.cpp‎

Lines changed: 37 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
#include <algorithm>
1515
#include <cmath>
1616
#include <cstddef>
17-
#include <exception>
1817
#include <functional>
1918
#include <future>
2019
#include <optional>
@@ -482,66 +481,53 @@ class parquet_field_struct_list : public parquet_field {
482481
auto const items_per_task = n / num_tasks;
483482
auto const remainder = n % num_tasks;
484483

485-
std::vector<std::future<void>> tasks;
486-
tasks.reserve(num_tasks);
487-
488484
// Size once up front, then fill by index below: workers read already-filled slots while the
489485
// main thread fills later ones, so a size-mutating emplace_back would be a data race.
490486
std::vector<cudf::host_span<uint8_t const>> all_ranges;
491487
all_ranges.resize(n);
492488

493-
// Wait for every task before all_ranges/val die on any exit path: a throw during collection
494-
// must not let a worker read freed memory (host-pool futures do not block on destruction).
495-
try {
496-
uint32_t struct_idx = 0;
497-
for (uint32_t task_id = 0; task_id < num_tasks; ++task_id) {
498-
auto const task_size = items_per_task + (task_id < remainder ? 1 : 0);
499-
auto const start_idx = std::min(n, struct_idx);
500-
auto const end_idx = std::min(n, start_idx + task_size);
501-
502-
if (start_idx >= end_idx) { break; }
503-
504-
// Collect struct ranges for the next task
505-
for (auto i = start_idx; i < end_idx; ++i) {
506-
uint8_t const* const start = cpr->m_cur;
507-
cpr->skip_struct_field(static_cast<int>(FieldType::STRUCT));
508-
all_ranges[i] =
509-
cudf::host_span<uint8_t const>{start, static_cast<size_t>(cpr->m_cur - start)};
510-
}
489+
// Declared after all_ranges so that every task is waited on before all_ranges dies on any
490+
// exit path (host-pool futures do not block on destruction)
491+
std::vector<std::future<void>> tasks;
492+
cudf::detail::futures_wait_guard tasks_guard{tasks};
493+
tasks.reserve(num_tasks);
511494

512-
// Launch task immediately to parse the structs in parallel while the main thread collects
513-
// the remaining ranges
514-
tasks.emplace_back(cudf::detail::host_worker_pool().submit_task(
515-
[&val = this->val, &all_ranges, start_idx, end_idx, mode]() {
516-
CompactProtocolReader local_cpr{nullptr, 0, mode};
517-
// end_idx alone bounds the loop: all_ranges[start_idx, end_idx) is fully populated
518-
// before this task launches, so no concurrent read of all_ranges.size() is needed.
519-
for (size_t i = start_idx; i < end_idx; ++i) {
520-
local_cpr.init(all_ranges[i].data(), all_ranges[i].size());
521-
local_cpr.read(&val[i]);
522-
}
523-
}));
524-
525-
struct_idx = end_idx;
526-
}
527-
} catch (...) {
528-
for (auto& task : tasks) {
529-
task.wait();
495+
uint32_t struct_idx = 0;
496+
for (uint32_t task_id = 0; task_id < num_tasks; ++task_id) {
497+
auto const task_size = items_per_task + (task_id < remainder ? 1 : 0);
498+
auto const start_idx = std::min(n, struct_idx);
499+
auto const end_idx = std::min(n, start_idx + task_size);
500+
501+
if (start_idx >= end_idx) { break; }
502+
503+
// Collect struct ranges for the next task
504+
for (auto i = start_idx; i < end_idx; ++i) {
505+
uint8_t const* const start = cpr->m_cur;
506+
cpr->skip_struct_field(static_cast<int>(FieldType::STRUCT));
507+
all_ranges[i] =
508+
cudf::host_span<uint8_t const>{start, static_cast<size_t>(cpr->m_cur - start)};
530509
}
531-
throw;
510+
511+
// Launch task immediately to parse the structs in parallel while the main thread collects
512+
// the remaining ranges
513+
tasks.emplace_back(cudf::detail::host_worker_pool().submit_task(
514+
[&val = this->val, &all_ranges, start_idx, end_idx, mode]() {
515+
CompactProtocolReader local_cpr{nullptr, 0, mode};
516+
// end_idx alone bounds the loop: all_ranges[start_idx, end_idx) is fully populated
517+
// before this task launches, so no concurrent read of all_ranges.size() is needed.
518+
for (size_t i = start_idx; i < end_idx; ++i) {
519+
local_cpr.init(all_ranges[i].data(), all_ranges[i].size());
520+
local_cpr.read(&val[i]);
521+
}
522+
}));
523+
524+
struct_idx = end_idx;
532525
}
533-
// Join EVERY task before leaving, then rethrow the first failure: task.get() throws on a
534-
// strict-mode mismatch, and bailing out of this loop mid-way would abandon the remaining
535-
// un-joined tasks, whose lambdas still reference all_ranges/val (use-after-free at exit).
536-
std::exception_ptr first_error;
526+
// Rethrows the first failure (e.g. a strict-mode mismatch); the guard waits for the
527+
// remaining tasks, whose lambdas reference all_ranges/val, while unwinding
537528
for (auto& task : tasks) {
538-
try {
539-
task.get();
540-
} catch (...) {
541-
if (!first_error) { first_error = std::current_exception(); }
542-
}
529+
task.get();
543530
}
544-
if (first_error) { std::rethrow_exception(first_error); }
545531
} else {
546532
// For small numbers of elements, use sequential processing to avoid overhead
547533
std::for_each(val.begin(), val.end(), [&cpr](auto& elem) { cpr->read(&elem); });

‎cpp/src/io/parquet/experimental/page_index_filter.cu‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -389,10 +389,11 @@ thrust::host_vector<bool> aggregate_reader_metadata::compute_data_page_mask(
389389
// task submission overheads
390390
auto constexpr max_tasks = 2;
391391
using task_page_row_offsets_type = std::vector<std::pair<std::vector<size_type>, size_type>>;
392-
std::vector<std::future<task_page_row_offsets_type>> page_row_offset_tasks{};
393-
page_row_offset_tasks.reserve(max_tasks);
394392
auto const cols_per_thread =
395393
cudf::detail::div_rounding_up_safe<std::size_t>(num_columns, max_tasks);
394+
std::vector<std::future<task_page_row_offsets_type>> page_row_offset_tasks{};
395+
cudf::detail::futures_wait_guard page_row_offset_tasks_guard{page_row_offset_tasks};
396+
page_row_offset_tasks.reserve(max_tasks);
396397

397398
// Submit page row offset compute tasks
398399
std::transform(cuda::counting_iterator<int>{0},

‎cpp/src/io/parquet/io_utils/parquet_io_utils.cpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ auto dispatch_tasks(std::size_t num_tasks, Task task)
106106
} else {
107107
// Dispatch the tasks to the host worker pool
108108
std::vector<std::future<result_type>> futures;
109+
cudf::detail::futures_wait_guard futures_guard{futures};
109110
futures.reserve(num_tasks);
110111
std::for_each(cuda::counting_iterator<std::size_t>(0),
111112
cuda::counting_iterator<std::size_t>(num_tasks),
@@ -237,6 +238,7 @@ void read_ranges_to_host(std::span<std::reference_wrapper<cudf::io::datasource>
237238
bool serialize_submissions)
238239
{
239240
std::vector<std::future<std::size_t>> host_read_tasks;
241+
cudf::detail::futures_wait_guard host_read_tasks_guard{host_read_tasks};
240242
std::vector<std::size_t> expected_sizes;
241243
host_read_tasks.reserve(count);
242244
expected_sizes.reserve(count);
@@ -383,7 +385,9 @@ fetch_byte_ranges_to_device_async_impl(
383385

384386
// Vectors to hold futures from datasource
385387
std::vector<std::future<size_t>> device_read_tasks{};
388+
cudf::detail::futures_wait_guard device_read_tasks_guard{device_read_tasks};
386389
std::vector<std::future<host_read_buffer>> host_read_tasks{};
390+
cudf::detail::futures_wait_guard host_read_tasks_guard{host_read_tasks};
387391
device_read_tasks.reserve(io_offsets.size());
388392
host_read_tasks.reserve(io_offsets.size());
389393

@@ -468,6 +472,7 @@ fetch_byte_ranges_to_device_async_impl(
468472
if (not host_buffers.empty()) { stream.sync(); }
469473

470474
auto sync_function = [](decltype(device_read_tasks) device_read_tasks) {
475+
cudf::detail::futures_wait_guard device_read_tasks_guard{device_read_tasks};
471476
for (auto& task : device_read_tasks) {
472477
task.get();
473478
}

‎cpp/src/io/parquet/reader_impl_helpers.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -645,6 +645,7 @@ void metadata::setup_page_index(std::span<uint8_t const> page_index_bytes, int64
645645
auto const remainder = total_column_chunks % num_tasks;
646646

647647
std::vector<std::future<void>> tasks;
648+
cudf::detail::futures_wait_guard tasks_guard{tasks};
648649
tasks.reserve(num_tasks);
649650

650651
std::size_t start_idx = 0;

‎cpp/src/io/parquet/reader_impl_helpers.hpp‎

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
#include <thrust/host_vector.h>
2020

2121
#include <algorithm>
22-
#include <exception>
2322
#include <functional>
2423
#include <future>
2524
#include <memory>
@@ -64,28 +63,19 @@ template <typename T, typename UnaryOp>
6463
}
6564

6665
std::vector<std::future<result_type>> tasks;
66+
cudf::detail::futures_wait_guard tasks_guard{tasks};
6767
tasks.reserve(inputs.size());
6868

69-
auto pending_exception = std::exception_ptr{};
70-
try {
71-
std::transform(inputs.begin(), inputs.end(), std::back_inserter(tasks), [&op](T const& input) {
72-
return cudf::detail::host_worker_pool().submit_task(
73-
[&op, input_ptr = &input] { return op(*input_ptr); });
74-
});
75-
} catch (...) {
76-
pending_exception = std::current_exception();
77-
}
69+
std::transform(inputs.begin(), inputs.end(), std::back_inserter(tasks), [&op](T const& input) {
70+
return cudf::detail::host_worker_pool().submit_task(
71+
[&op, input_ptr = &input] { return op(*input_ptr); });
72+
});
7873

74+
// Rethrows the first task failure; the guard waits for the remaining tasks on unwinding
7975
for (auto& task : tasks) {
80-
try {
81-
results.emplace_back(task.get());
82-
} catch (...) {
83-
if (not pending_exception) { pending_exception = std::current_exception(); }
84-
}
76+
results.emplace_back(task.get());
8577
}
8678

87-
if (pending_exception) { std::rethrow_exception(pending_exception); }
88-
8979
return results;
9080
}
9181

0 commit comments

Comments
 (0)