defaults.hpp

/*
 * Copyright (c) 2022-2025, 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 <cstddef>
#include <cstdlib>
#include <initializer_list>
#include <sstream>
#include <stdexcept>
#include <string>
#include <type_traits>
 
#include <kvikio/compat_mode.hpp>
#include <kvikio/error.hpp>
#include <kvikio/http_status_codes.hpp>
#include <kvikio/shim/cufile.hpp>
#include <kvikio/threadpool_wrapper.hpp>
 
namespace kvikio {
 
template <typename T>
T getenv_or(std::string_view env_var_name, T default_val)
{
  auto const* 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 constexpr (!std::is_same_v<T, std::string>) {
    KVIKIO_EXPECT(!sstream.fail(),
                  "unknown config value " + std::string{env_var_name} + "=" + std::string{env_val},
                  std::invalid_argument);
  }
 
  return converted_val;
}
 
template <>
bool getenv_or(std::string_view env_var_name, bool default_val);
 
template <>
CompatMode getenv_or(std::string_view env_var_name, CompatMode default_val);
 
template <>
std::vector<int> getenv_or(std::string_view env_var_name, std::vector<int> default_val);
 
template <typename T>
std::tuple<std::string_view, T, bool> getenv_or(
  std::initializer_list<std::string_view> env_var_names, T default_val)
{
  KVIKIO_EXPECT(env_var_names.size() > 0,
                "`env_var_names` must contain at least one environment variable name.",
                std::invalid_argument);
  std::string_view env_name_target;
  std::string_view env_val_target;
 
  for (auto const& env_var_name : env_var_names) {
    auto const* env_val = std::getenv(env_var_name.data());
    if (env_val == nullptr) { continue; }
 
    if (!env_name_target.empty() && env_val_target != env_val) {
      std::stringstream ss;
      ss << "Environment variable " << env_var_name << " (" << env_val
         << ") has already been set by its alias " << env_name_target << " (" << env_val_target
         << ") with a different value.";
      KVIKIO_FAIL(ss.str(), std::invalid_argument);
    }
 
    env_name_target = env_var_name;
    env_val_target  = env_val;
  }
 
  if (env_name_target.empty()) { return {env_name_target, default_val, false}; }
 
  auto res = getenv_or<T>(env_name_target, default_val);
  return {env_name_target, res, true};
}
 
class defaults {
 private:
  BS_thread_pool _thread_pool{get_num_threads_from_env()};
  CompatMode _compat_mode;
  std::size_t _task_size;
  std::size_t _gds_threshold;
  std::size_t _bounce_buffer_size;
  std::size_t _http_max_attempts;
  long _http_timeout;
  std::vector<int> _http_status_codes;
 
  static unsigned int get_num_threads_from_env();
 
  defaults();
 
  KVIKIO_EXPORT static defaults* instance();
 
 public:
  [[nodiscard]] static CompatMode compat_mode();
 
  static void set_compat_mode(CompatMode compat_mode);
 
  static CompatMode infer_compat_mode_if_auto(CompatMode compat_mode) noexcept;
 
  static bool is_compat_mode_preferred(CompatMode compat_mode) noexcept;
 
  static bool is_compat_mode_preferred();
 
  [[nodiscard]] static BS_thread_pool& thread_pool();
 
  [[nodiscard]] static unsigned int thread_pool_nthreads();
 
  static void set_thread_pool_nthreads(unsigned int nthreads);
 
  [[nodiscard]] static unsigned int num_threads();
 
  static void set_num_threads(unsigned int nthreads);
 
  [[nodiscard]] static std::size_t task_size();
 
  static void set_task_size(std::size_t nbytes);
 
  [[nodiscard]] static std::size_t gds_threshold();
 
  static void set_gds_threshold(std::size_t nbytes);
 
  [[nodiscard]] static std::size_t bounce_buffer_size();
 
  static void set_bounce_buffer_size(std::size_t nbytes);
 
  [[nodiscard]] static std::size_t http_max_attempts();
 
  static void set_http_max_attempts(std::size_t attempts);
 
  [[nodiscard]] static long http_timeout();
 
  static void set_http_timeout(long timeout_seconds);
 
  [[nodiscard]] static std::vector<int> const& http_status_codes();
 
  static void set_http_status_codes(std::vector<int> status_codes);
};
 
}  // namespace kvikio