Improve design a bit

rapidsai · Oct 20, 2024 · df9f756 · df9f756
1 parent c230b70
commit df9f756
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 10 deletions.
diff --git 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
@@ -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<unsigned int>("KVIKIO_NTHREADS", 8U); }
+  auto nthreads = getenv_or<unsigned int>("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
@@ -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
@@ -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");