From 1683e3a9e2059e8330390c93d276ff784eaf5045 Mon Sep 17 00:00:00 2001 From: Timothy Kelley Date: Mon, 10 Jun 2019 12:15:38 -0600 Subject: [PATCH] Add function to query environment vars; class to represent SLURM information. Improve how get_env_val communicates whether a key is defined in the environment. Include . Correct warnable syntactic shortcoming. Fix directory mismatch in doxygen comments. Fix doxygen usage. Another doxygen fix. Another clang-format. --- src/c4/QueryEnv.hh | 118 +++++++++++++++++++++++++ src/c4/test/tstQueryEnv.cc | 170 +++++++++++++++++++++++++++++++++++++ 2 files changed, 288 insertions(+) create mode 100644 src/c4/QueryEnv.hh create mode 100644 src/c4/test/tstQueryEnv.cc diff --git a/src/c4/QueryEnv.hh b/src/c4/QueryEnv.hh new file mode 100644 index 0000000000..f01caefb99 --- /dev/null +++ b/src/c4/QueryEnv.hh @@ -0,0 +1,118 @@ +//----------------------------------*-C++-*----------------------------------// +/*! + * \file c4/QueryEnv.hh + * \author Tim Kelley + * \date Fri Jun 7 08:06:53 2019 + * \brief Functions for working with your environment + * \note Copyright (C) 2019 Triad National Security, LLC. + * All rights reserved. */ +//---------------------------------------------------------------------------// + +#ifndef Query_Env_hh +#define Query_Env_hh + +#include // getenv, setenv +#include +#include +#include +#include +#include +#include // pair + +namespace rtt_c4 { + +//----------------------------------------------------------------------------// +/*! + * \brief Get a value from the environment + * \tparam T: the type of the value, must be default constructible + * \param key: the key to which you would like the corresponding value + * \param default_value: a default value to return if key undefined. + * \return: {whether key was defined, corresponding value} + * + * \note Calls POSIX getenv() function, which is not required to be re-entrant. + * If the key was set in the environment, get_env_val() returns the value + * converted to type T, that was set in the environment. If the key was not + * set in the environment, it returns the default_value that the caller + * provides. The first argument of the pair describes whether the key was defined. + */ +template +std::pair get_env_val(std::string const &key, T default_value = T{}) { + static_assert(std::is_default_constructible::value, + "get_env_val only usable with default constructible types"); + T val{default_value}; + char *s_val = getenv(key.c_str()); + bool is_defined(false); + if (s_val) { + std::stringstream val_str(s_val); + val_str >> val; + is_defined = true; + } + return std::make_pair(is_defined, val); +} // get_env_val + +//----------------------------------------------------------------------------// +/*! + *\brief Basic information about SLURM tasks, and whether that information + * was available. + * + * \note: values are based on what was in the environment at the time this + * class is instantiated. This class is likely non-reentrant, so use with care + * in multithreaded environments. This class relies on SLURM setting the + * following environment variables: + * + * SLURM_CPUS_PER_TASK (the argument to -c) + * SLURM_NTASKS (the argument to -n) + * SLURM_JOB_NUM_NODES (the argument to -N) + * + * Draco's unit tests don't really make sure that's the case, so if SLURM + * changes, this may break. */ +class SLURM_Task_Info { +public: + /**\brief Get SLURM_CPUS_PER_TASK */ + int get_cpus_per_task() const { return cpus_per_task_; } + + /**\brief Get SLURM_NTASKS */ + int get_ntasks() const { return ntasks_; } + + /**\brief Get SLURM_JOB_NUM_NODES */ + int get_job_num_nodes() const { return job_num_nodes_; } + + /* note: these rely on the idea that n_cpus_per_task etc are never + * going to be in the realm of 2 billion. On the blessaed day that + * comes to pass, Machine Overlords, rethink this (please / thank you). */ + + /**\brief Was SLURM_CPUS_PER_TASK set? */ + bool is_cpus_per_task_set() const { return def_cpus_per_task_; } + + /**\brief Was SLURM_NTASKS set? */ + bool is_ntasks_set() const { return def_ntasks_; } + + /**\brief Was SLURM_JOB_NUM_NODES set? */ + bool is_job_num_nodes_set() const { return def_job_num_nodes_; } + + // ctor + SLURM_Task_Info() { + std::tie(def_cpus_per_task_, cpus_per_task_) = + get_env_val("SLURM_CPUS_PER_TASK", cpus_per_task_); + std::tie(def_ntasks_, ntasks_) = get_env_val("SLURM_NTASKS", ntasks_); + std::tie(def_job_num_nodes_, job_num_nodes_) = + get_env_val("SLURM_JOB_NUM_NODES", job_num_nodes_); + } + + // state +private: + int cpus_per_task_{0xFFFFFFF}; //!< arg to -c + int ntasks_{0xFFFFFFE}; //!< arg to -n + int job_num_nodes_{0xFFFFFFD}; //!< arg to -N + bool def_cpus_per_task_{false}; //!< whether SLURM_CPUS_PER_TASK was defined + bool def_ntasks_{false}; //!< whether SLURM_NTASKS was defined + bool def_job_num_nodes_{false}; //!< whether SLURM_JOB_NUM_NODES was defined +}; // SLURM_Task_Info + +} // namespace rtt_c4 + +#endif // Query_Env_hh + +//---------------------------------------------------------------------------// +// end of QueryEnv.hh +//---------------------------------------------------------------------------// diff --git a/src/c4/test/tstQueryEnv.cc b/src/c4/test/tstQueryEnv.cc new file mode 100644 index 0000000000..88c2c1996b --- /dev/null +++ b/src/c4/test/tstQueryEnv.cc @@ -0,0 +1,170 @@ +//----------------------------------*-C++-*----------------------------------// +/*! + * \file c4/test/tstQueryEnv.cc + * \author Tim Kelley + * \date Fri Jun 7 08:06:53 2019 + * \note Copyright (C) 2019 Triad National Security, LLC. + * All rights reserved. */ +//---------------------------------------------------------------------------// + +#include "c4/QueryEnv.hh" +#include "ds++/Release.hh" +#include "ds++/ScalarUnitTest.hh" +#include +#include // std::function +#include +#include + +using rtt_dsxx::UnitTest; + +using env_store_value = std::pair; +using env_store_t = std::map; + +//----------------------------------------------------------------------------// +/* Helper function: Record SLURM keys and values, if any, then remove them from + * environment. Return recorded values so they can be restored later. */ +env_store_t clean_env() { + // for each key, is it defined? If so, record the value, then unset it. If + // not, note that. + std::string slurm_keys[] = {"SLURM_CPUS_PER_TASK", "SLURM_NTASKS", + "SLURM_JOB_NUM_NODES"}; + env_store_t store{}; + for (auto &k : slurm_keys) { + // We're assuming none of those flags were set to something like 2 billion + constexpr int a_large_int = 0x0FFFFFFC; + bool k_defined{false}; + int ival; + std::tie(k_defined, ival) = rtt_c4::get_env_val(k, a_large_int); + if (k_defined) { + std::stringstream storestr; + storestr << ival; + store.insert({k, {true, storestr.str()}}); + int unset_ok = unsetenv(k.c_str()); + if (0 != unset_ok) { + printf("%s:%i Failed to unset env var %s! errno = %d\n", __FUNCTION__, + __LINE__, k.c_str(), errno); + // throw something? + } + } else { + store.insert({k, {false, ""}}); + } + } // for k in slurm keys + return store; +} // clean_env + +//----------------------------------------------------------------------------// +/* Helper function: Restore the SLURM keys that were previously defined. */ +void restore_env(env_store_t const &store) { + /* For each key, if it was defined, restore that definition. If it was + * not defined, destroy any subsequent definition */ + for (auto str_it : store) { + std::string const &key{str_it.first}; + env_store_value const &val{str_it.second}; + bool const &was_defined{val.first}; + if (was_defined) { + std::string const &val_str{val.second}; + int overwrite{1}; // yes, overwrite existing values + int set_ok = setenv(key.c_str(), val_str.c_str(), overwrite); + if (0 != set_ok) { + printf("%s:%i Failed to set env var %s to %s, errno = %d\n", + __FUNCTION__, __LINE__, key.c_str(), val_str.c_str(), errno); + // throw something + } + } else { + int unset_ok = unsetenv(key.c_str()); + if (0 != unset_ok) { + printf("%s:%i Failed to unset env var %s! errno = %d\n", __FUNCTION__, + __LINE__, key.c_str(), errno); + // throw something? + } + } + } // for things in store + return; +} // restore_env + +//----------------------------------------------------------------------------// +/* Test with a "clean" environment--that is, no slurm keys. */ +void test_instantiate_SLURM_Info(UnitTest &ut) { + /* Test instantiating SLURM_Task_Info in a 'clean' environment */ + auto env_tmp = clean_env(); + rtt_c4::SLURM_Task_Info ti; + FAIL_IF(ti.is_cpus_per_task_set()); + FAIL_IF(ti.is_ntasks_set()); + FAIL_IF(ti.is_job_num_nodes_set()); + FAIL_IF_NOT(ti.get_cpus_per_task() == 0xFFFFFFF); + FAIL_IF_NOT(ti.get_ntasks() == 0xFFFFFFE); + FAIL_IF_NOT(ti.get_job_num_nodes() == 0xFFFFFFD); + restore_env(env_tmp); + return; +} + +//----------------------------------------------------------------------------// +/* Test with a "live" environment--that is, slurm keys are defined. */ +void test_SLURM_Info(UnitTest &ut) { + /* Test instantiating SLURM_Task_Info in a 'clean' environment */ + auto orig_env = clean_env(); + int const iset_cpus_per_task{21}, iset_ntasks{341}, iset_job_num_nodes{1001}; + const char *set_cpus_per_task{"21"}, *set_ntasks{"341"}, + *set_job_num_nodes{"1001"}; + env_store_t test_env = {{"SLURM_CPUS_PER_TASK", {true, set_cpus_per_task}}, + {"SLURM_NTASKS", {true, set_ntasks}}, + {"SLURM_JOB_NUM_NODES", {true, set_job_num_nodes}}}; + restore_env(test_env); + rtt_c4::SLURM_Task_Info ti; + FAIL_IF_NOT(ti.is_cpus_per_task_set()); + FAIL_IF_NOT(ti.is_ntasks_set()); + FAIL_IF_NOT(ti.is_job_num_nodes_set()); + FAIL_IF_NOT(ti.get_cpus_per_task() == iset_cpus_per_task); + FAIL_IF_NOT(ti.get_ntasks() == iset_ntasks); + FAIL_IF_NOT(ti.get_job_num_nodes() == iset_job_num_nodes); + restore_env(orig_env); + return; +} + +//----------------------------------------------------------------------------// +/* Test with a partial "live" environment--that is, slurm keys are defined. */ +void test_SLURM_Info_partial(UnitTest &ut) { + /* Test instantiating SLURM_Task_Info in a 'clean' environment */ + auto orig_env = clean_env(); + int const iset_cpus_per_task{21}, iset_job_num_nodes{1001}; + const char *set_cpus_per_task{"21"}, *set_job_num_nodes{"1001"}; + env_store_t test_env = {{"SLURM_CPUS_PER_TASK", {true, set_cpus_per_task}}, + {"SLURM_JOB_NUM_NODES", {true, set_job_num_nodes}}}; + restore_env(test_env); + rtt_c4::SLURM_Task_Info ti; + FAIL_IF_NOT(ti.is_cpus_per_task_set()); + FAIL_IF(ti.is_ntasks_set()); + FAIL_IF_NOT(ti.is_job_num_nodes_set()); + FAIL_IF_NOT(ti.get_cpus_per_task() == iset_cpus_per_task); + FAIL_IF_NOT(ti.get_ntasks() == 0xFFFFFFE); + FAIL_IF_NOT(ti.get_job_num_nodes() == iset_job_num_nodes); + restore_env(orig_env); + return; +} + +using t_func = std::function; + +//----------------------------------------------------------------------------// +void run_a_test(UnitTest &u, t_func f, std::string const &msg) { + f(u); + if (u.numFails == 0) { + u.passes(msg); + } + return; +} + +//----------------------------------------------------------------------------// +int main(int argc, char *argv[]) { + rtt_dsxx::ScalarUnitTest ut(argc, argv, rtt_dsxx::release); + try { + run_a_test(ut, test_instantiate_SLURM_Info, "SLURM_Info (clean env) ok."); + run_a_test(ut, test_SLURM_Info, "SLURM_Info (SLURM vars set) ok."); + run_a_test(ut, test_SLURM_Info_partial, + "SLURM_Info (partial SLURM vars set) ok."); + } // try--catches in the epilog: + UT_EPILOG(ut); +} + +//---------------------------------------------------------------------------// +// end of c4/test/tstQueryEnv.cc +//---------------------------------------------------------------------------//