Merge pull request #634 from timmah/feat-add-slurm-vars

Add function to query environment vars; class to represent SLURM info…

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 <cstdlib> // getenv, setenv
+#include <iostream>
+#include <sstream>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <utility> // 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 <typename T>
+std::pair<bool, T> get_env_val(std::string const &key, T default_value = T{}) {
+  static_assert(std::is_default_constructible<T>::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<int>("SLURM_CPUS_PER_TASK", cpus_per_task_);
+    std::tie(def_ntasks_, ntasks_) = get_env_val<int>("SLURM_NTASKS", ntasks_);
+    std::tie(def_job_num_nodes_, job_num_nodes_) =
+        get_env_val<int>("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
+//---------------------------------------------------------------------------//

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 <cstdlib>
+#include <functional> // std::function
+#include <map>
+#include <thread>
+
+using rtt_dsxx::UnitTest;
+
+using env_store_value = std::pair<bool, std::string>;
+using env_store_t = std::map<std::string, env_store_value>;
+
+//----------------------------------------------------------------------------//
+/* 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<int>(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(UnitTest &)>;
+
+//----------------------------------------------------------------------------//
+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
+//---------------------------------------------------------------------------//