defaults.hpp

/*
 * Copyright (c) 2022-2023, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <algorithm>
#include <cstddef>
#include <cstdlib>
#include <sstream>
#include <stdexcept>
#include <string>
#include <utility>
 
#include <kvikio/shim/cufile.hpp>
#include <kvikio/thread_pool.hpp>
 
namespace kvikio {
namespace detail {
 
template <typename T>
T getenv_or(std::string_view env_var_name, T default_val)
{
  const auto* env_val = std::getenv(env_var_name.data());
  if (env_val == nullptr) { return default_val; }
 
  std::stringstream sstream(env_val);
  T converted_val;
  sstream >> converted_val;
  if (sstream.fail()) {
    throw std::invalid_argument("unknown config value " + std::string{env_var_name} + "=" +
                                std::string{env_val});
  }
  return converted_val;
}
 
template <>
inline bool getenv_or(std::string_view env_var_name, bool default_val)
{
  const auto* env_val = std::getenv(env_var_name.data());
  if (env_val == nullptr) { return default_val; }
  try {
    // Try parsing `env_var_name` as a integer
    return static_cast<bool>(std::stoi(env_val));
  } catch (const std::invalid_argument&) {
  }
  // Convert to lowercase
  std::string str{env_val};
  std::transform(str.begin(), str.end(), str.begin(), ::tolower);
  // Trim whitespaces
  std::stringstream trimmer;
  trimmer << str;
  str.clear();
  trimmer >> str;
  // Match value
  if (str == "true" || str == "on" || str == "yes") { return true; }
  if (str == "false" || str == "off" || str == "no") { return false; }
  throw std::invalid_argument("unknown config value " + std::string{env_var_name} + "=" +
                              std::string{env_val});
}
 
}  // namespace detail
 
class defaults {
 private:
  kvikio::third_party::thread_pool _thread_pool{get_num_threads_from_env()};
  bool _compat_mode;
  std::size_t _task_size;
  std::size_t _gds_threshold;
 
  static unsigned int get_num_threads_from_env()
  {
    const int ret = detail::getenv_or("KVIKIO_NTHREADS", 1);
    if (ret <= 0) { throw std::invalid_argument("KVIKIO_NTHREADS has to be a positive integer"); }
    return ret;
  }
 
  defaults()
  {
    // Determine the default value of `compat_mode`
    {
      if (std::getenv("KVIKIO_COMPAT_MODE") != nullptr) {
        // Setting `KVIKIO_COMPAT_MODE` take precedence
        _compat_mode = detail::getenv_or("KVIKIO_COMPAT_MODE", false);
      } else {
        // If `KVIKIO_COMPAT_MODE` isn't set, we infer based on runtime environment
        _compat_mode = !is_cufile_available();
      }
    }
    // Determine the default value of `task_size`
    {
      const ssize_t env = detail::getenv_or("KVIKIO_TASK_SIZE", 4 * 1024 * 1024);
      if (env <= 0) {
        throw std::invalid_argument("KVIKIO_TASK_SIZE has to be a positive integer");
      }
      _task_size = env;
    }
    // Determine the default value of `gds_threshold`
    {
      const ssize_t env = detail::getenv_or("KVIKIO_GDS_THRESHOLD", 1024 * 1024);
      if (env <= 0) {
        throw std::invalid_argument("KVIKIO_GDS_THRESHOLD has to be a positive integer");
      }
      _gds_threshold = env;
    }
  }
 
  static defaults* instance()
  {
    static defaults _instance;
    return &_instance;
  }
 
 public:
  [[nodiscard]] static bool compat_mode() { return instance()->_compat_mode; }
 
  static void compat_mode_reset(bool enable) { instance()->_compat_mode = enable; }
 
  [[nodiscard]] static kvikio::third_party::thread_pool& thread_pool()
  {
    return instance()->_thread_pool;
  }
 
  [[nodiscard]] static unsigned int thread_pool_nthreads()
  {
    return thread_pool().get_thread_count();
  }
 
  static void thread_pool_nthreads_reset(unsigned int nthreads) { thread_pool().reset(nthreads); }
 
  [[nodiscard]] static std::size_t task_size() { return instance()->_task_size; }
 
  static void task_size_reset(std::size_t nbytes) { instance()->_task_size = nbytes; }
 
  [[nodiscard]] static std::size_t gds_threshold() { return instance()->_gds_threshold; }
 
  static void gds_threshold_reset(std::size_t nbytes) { instance()->_gds_threshold = nbytes; }
};
 
}  // namespace kvikio