diff --git a/cpp/include/kvikio/detail/parallel_operation.hpp b/cpp/include/kvikio/detail/parallel_operation.hpp index 0e0ef3cdee..53403e80a4 100644 --- a/cpp/include/kvikio/detail/parallel_operation.hpp +++ b/cpp/include/kvikio/detail/parallel_operation.hpp @@ -64,6 +64,24 @@ inline const std::pair get_next_color_and return {nvtx_color, call_idx}; } +/** + * @brief Whether the sub-range tasks of one parallel operation join without blocking a worker. Set + * by the environment variable `KVIKIO_NONBLOCKING_TASK_JOIN`. + * + * - Disabled (default): the last task performs its own sub-range and then blocks on the futures of + * its siblings. One worker therefore block-waits for the duration of the join, and the operation + * completes on that worker. + * - Enabled: every sub-range is submitted as an independent task, and whichever one finishes last + * completes the operation. No worker block-waits. + * + * @return True if the non-blocking join is enabled. + */ +inline bool nonblocking_task_join() +{ + static bool const enabled = getenv_or("KVIKIO_NONBLOCKING_TASK_JOIN", false); + return enabled; +} + /** * @brief Options for a single I/O task submission. */ diff --git a/cpp/src/remote_handle.cpp b/cpp/src/remote_handle.cpp index be8da42a13..8c57f9c487 100644 --- a/cpp/src/remote_handle.cpp +++ b/cpp/src/remote_handle.cpp @@ -793,6 +793,69 @@ namespace { : IoBackend::REMOTE_HTTP; } +/** + * @brief Submit each `task_size` sub-range of a read to `thread_pool` as an independent task. + * + * The task that finishes last resolves the returned future, and no worker is block-waiting on its + * siblings. + * + * @tparam Task Callable that performs one sub-range read. Must be copy-constructible. + * @param task Performs one sub-range read. + * @param buf Destination buffer. + * @param size Number of bytes to read. Must be positive. + * @param file_offset File offset in bytes. + * @param task_size Size of each sub-range in bytes. Must be positive. + * @param thread_pool Thread pool to run the tasks on. Must not be null. + * @param recorder Records the logical operation. Null when nobody is observing. + * @return Future that on completion returns the number of bytes read, which is always `size`. + */ +template +std::future submit_subrange_reads( + Task task, + std::byte* buf, + std::size_t size, + std::size_t file_offset, + std::size_t task_size, + ThreadPool* thread_pool, + std::shared_ptr recorder) +{ + KVIKIO_NVTX_FUNC_RANGE(size); + + std::size_t const num_subranges = (task_size >= size) ? 1 : (size + task_size - 1) / task_size; + auto aggregate = std::make_shared(num_subranges); + aggregate->recorder = std::move(recorder); + auto fut = aggregate->get_future(); + + auto const& [nvtx_color, call_idx] = detail::get_next_color_and_call_idx(); + + std::size_t remaining = size; + std::size_t cur_off = file_offset; + auto cur_buf = buf; + for (std::size_t i = 0; i < num_subranges; ++i) { + std::size_t const subrange_size = std::min(task_size, remaining); + thread_pool->detach_task([task, + aggregate, + cur_buf, + subrange_size, + cur_off, + call_idx = call_idx, + nvtx_color = nvtx_color] { + KVIKIO_NVTX_SCOPED_RANGE("task", call_idx, nvtx_color); + try { + auto const nbytes = task(cur_buf, subrange_size, cur_off); + aggregate->on_subrange_complete(nbytes); + } catch (...) { + aggregate->on_subrange_failed(std::current_exception()); + } + }); + cur_buf += subrange_size; + cur_off += subrange_size; + remaining -= subrange_size; + } + + return fut; +} + } // namespace std::size_t RemoteHandle::read(void* buf, std::size_t size, std::size_t file_offset) @@ -916,6 +979,20 @@ std::future RemoteHandle::pread(void* buf, : nullptr; if (io_backend == RemoteIOBackend::EASY_THREADPOOL) { + if (detail::nonblocking_task_join()) { + auto task = [this, is_host_mem]( + std::byte* dst, std::size_t size, std::size_t file_offset) -> std::size_t { + return read_impl(dst, size, file_offset, is_host_mem); + }; + return submit_subrange_reads(std::move(task), + static_cast(buf), + size, + file_offset, + task_size, + thread_pool, + std::move(recorder)); + } + auto& [nvtx_color, call_idx] = detail::get_next_color_and_call_idx(); auto task = [this, is_host_mem](void* devPtr_base, @@ -934,7 +1011,7 @@ std::future RemoteHandle::pread(void* buf, {.thread_pool = thread_pool, .call_idx = call_idx, .nvtx_color = nvtx_color, - .recorder = recorder}); + .recorder = std::move(recorder)}); } // MULTI_POLL path. The lifecycle of one pread() call uses four cooperating pieces: diff --git a/docs/source/runtime_settings.rst b/docs/source/runtime_settings.rst index 5d229855e9..f8563c5b36 100644 --- a/docs/source/runtime_settings.rst +++ b/docs/source/runtime_settings.rst @@ -33,6 +33,15 @@ When opportunistic Direct I/O read is enabled (``KVIKIO_AUTO_DIRECT_IO_READ=1``) This setting can be queried (:py:func:`kvikio.defaults.get`) and modified (:py:func:`kvikio.defaults.set`) at runtime using the property name ``task_size``. +Non-blocking Task Join ``KVIKIO_NONBLOCKING_TASK_JOIN`` +-------------------------------------------------------- +Controls how the tasks of one parallel operation are joined once they have been submitted to the thread pool. Set the environment variable ``KVIKIO_NONBLOCKING_TASK_JOIN`` to ``true``, ``on``, ``yes``, or ``1`` (case-insensitive) to enable the non-blocking join. It is disabled by default. + + * Disabled (default): the last task performs its own sub-range and then blocks on the futures of its siblings. One worker therefore block-waits for the duration of the join, and the operation completes on that worker. + * Enabled: every sub-range is submitted as an independent task, and whichever one finishes last completes the operation. No worker block-waits. + +The variable is read only from the environment, and only once on first use. It currently applies to remote (HTTP/S3/WebHDFS) reads under the ``EASY_THREADPOOL`` backend, and has no effect under ``MULTI_POLL``. + GDS Threshold ``KVIKIO_GDS_THRESHOLD`` -------------------------------------- In order to improve performance of small IO, ``.pread()`` and ``.pwrite()`` implement a shortcut that circumvent the threadpool and use the POSIX backend directly. Set the environment variable ``KVIKIO_GDS_THRESHOLD`` to the minimum size (in bytes) to use GDS. If not set, the default value is 16384 (16 KiB).