Address reviewer comments

rapidsai · Oct 22, 2024 · eea4bdf · eea4bdf
1 parent a739292
commit eea4bdf
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 21 deletions.
diff --git a/cpp/include/cudf/io/config_utils.hpp b/cpp/include/cudf/io/config_utils.hpp
@@ -36,6 +36,12 @@ bool is_gds_enabled();
  */
 bool is_kvikio_enabled();
 
+/**
+ * @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_from_env();
+
 }  // namespace cufile_integration
 
 namespace nvcomp_integration {
@@ -51,15 +57,5 @@ bool is_all_enabled();
 bool is_stable_enabled();
 
 }  // namespace nvcomp_integration
-
-namespace kvikio_setting {
-
-/**
- * @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_from_env();
-
-}  // namespace kvikio_setting
 }  // namespace io
 }  // namespace CUDF_EXPORT cudf
diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp
@@ -22,6 +22,7 @@
 #include <kvikio/defaults.hpp>
 
 #include <cstdlib>
+#include <mutex>
 #include <sstream>
 #include <string>
 
@@ -55,6 +56,14 @@ bool is_gds_enabled() { return is_always_enabled() or get_env_policy() == usage_
 
 bool is_kvikio_enabled() { return get_env_policy() == usage_policy::KVIKIO; }
 
+void set_thread_pool_nthreads_from_env()
+{
+  static std::once_flag flag{};
+  std::call_once(flag, [] {
+    auto nthreads = getenv_or<unsigned int>("KVIKIO_NTHREADS", 8U);
+    kvikio::defaults::thread_pool_nthreads_reset(nthreads);
+  });
+}
 }  // namespace cufile_integration
 
 namespace nvcomp_integration {
@@ -83,13 +92,4 @@ bool is_all_enabled() { return get_env_policy() == usage_policy::ALWAYS; }
 bool is_stable_enabled() { return is_all_enabled() or get_env_policy() == usage_policy::STABLE; }
 
 }  // namespace nvcomp_integration
-
-namespace kvikio_setting {
-void set_thread_pool_nthreads_from_env()
-{
-  auto nthreads = getenv_or<unsigned int>("KVIKIO_NTHREADS", 8U);
-  kvikio::defaults::thread_pool_nthreads_reset(nthreads);
-}
-}  // namespace kvikio_setting
-
 }  // namespace cudf::io
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_from_env();
+      cufile_integration::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_from_env();
+      cufile_integration::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");