diff --git a/cpp/include/cudf/io/config_utils.hpp b/cpp/include/cudf/io/config_utils.hpp index f4ac573fed4..b4ac4fdd155 100644 --- a/cpp/include/cudf/io/config_utils.hpp +++ b/cpp/include/cudf/io/config_utils.hpp @@ -55,13 +55,10 @@ bool is_stable_enabled(); namespace kvikio_setting { /** - * @brief Set kvikIO thread pool size. - * - * @param nthreads The number of threads used for kvikIO thread pool. If not set, the environment - * variable KVIKIO_NTHREADS will be used. If KVIKIO_NTHREADS is unset either, use 8 threads by - * default. Priority: nthreads parameter > KVIKIO_NTHREADS > default 8 threads. + * @brief Set kvikIO thread pool size according to the environment variable KVIKIO_NTHREADS. If + * KVIKIO_NTHREADS is not set, use 8 threads by default. */ -void set_thread_pool_nthreads(unsigned int nthreads = 0U); +void set_thread_pool_nthreads_from_env(); } // namespace kvikio_setting } // namespace io diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index 50fa21b44ea..1ccedac8cf2 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -85,9 +85,9 @@ bool is_stable_enabled() { return is_all_enabled() or get_env_policy() == usage_ } // namespace nvcomp_integration namespace kvikio_setting { -void set_thread_pool_nthreads(unsigned int nthreads) +void set_thread_pool_nthreads_from_env() { - if (nthreads == 0U) { nthreads = getenv_or("KVIKIO_NTHREADS", 8U); } + auto nthreads = getenv_or("KVIKIO_NTHREADS", 8U); kvikio::defaults::thread_pool_nthreads_reset(nthreads); } } // namespace kvikio_setting diff --git a/cpp/src/io/utilities/data_sink.cpp b/cpp/src/io/utilities/data_sink.cpp index 8274ff3fa63..2e838922747 100644 --- a/cpp/src/io/utilities/data_sink.cpp +++ b/cpp/src/io/utilities/data_sink.cpp @@ -42,7 +42,7 @@ class file_sink : public data_sink { if (!_output_stream.is_open()) { detail::throw_on_file_open_failure(filepath, true); } if (cufile_integration::is_kvikio_enabled()) { - kvikio_setting::set_thread_pool_nthreads(); + kvikio_setting::set_thread_pool_nthreads_from_env(); _kvikio_file = kvikio::FileHandle(filepath, "w"); CUDF_LOG_INFO("Writing a file using kvikIO, with compatibility mode {}.", _kvikio_file.is_compat_mode_on() ? "on" : "off"); diff --git a/cpp/src/io/utilities/datasource.cpp b/cpp/src/io/utilities/datasource.cpp index cbef30b6034..7448abfd313 100644 --- a/cpp/src/io/utilities/datasource.cpp +++ b/cpp/src/io/utilities/datasource.cpp @@ -48,7 +48,7 @@ class file_source : public datasource { { detail::force_init_cuda_context(); if (cufile_integration::is_kvikio_enabled()) { - kvikio_setting::set_thread_pool_nthreads(); + kvikio_setting::set_thread_pool_nthreads_from_env(); _kvikio_file = kvikio::FileHandle(filepath); CUDF_LOG_INFO("Reading a file using kvikIO, with compatibility mode {}.", _kvikio_file.is_compat_mode_on() ? "on" : "off");