Implement monitor disable via environment (#183)

Allow monitors to be disabled via the PRMON_DISABLE_MONITOR
environment variable

README.md

@@ -108,6 +108,20 @@ incomplete arguments. If `prmon` starts a program itself (using `--`) then
 When invoked with `-h` or `--help` usage information is printed, as well as a
 list of all available monitoring components.
 
+### Environment Variables
+
+The `PRMON_DISABLE_MONITOR` environment variable can be used to specifiy a comma
+separated list of monitor names that will be disabled. This is useful when
+`prmon` is being invoked by some other part of a job or workflow, so the user
+does not have direct access to the command line options used. e.g.
+
+```sh
+export PRMON_DISABLE_MONITOR=nvidiamon
+other_code_that_invokes_prmon
+...
+```
+
+Disables the `nvidiamon` monitor.
 ## Outputs
 
 In the `filename` output file, plain text with statistics written every

package/src/prmon.cpp

@@ -289,6 +289,9 @@ int main(int argc, char* argv[]) {
     }
   }
 
+  // Get additional configuration from the environment
+  prmon::disable_monitors_from_env(disabled_monitors);
+
   if (do_help) {
     std::cout
         << "prmon is a process monitor program that records runtime data\n"

package/src/prmonutils.cpp

@@ -7,6 +7,7 @@
 #include <sys/stat.h>
 
 #include <cstddef>
+#include <cstdlib>
 #include <deque>
 #include <fstream>
 #include <iostream>
@@ -38,7 +39,7 @@ std::vector<pid_t> pstree_pids(const pid_t mother_pid) {
   char smaps_buffer[64];
   snprintf(smaps_buffer, 64, "pstree -A -p %ld | tr \\- \\\\n",
            (long)mother_pid);
-  FILE* pipe = popen(smaps_buffer, "r");
+  FILE *pipe = popen(smaps_buffer, "r");
   if (pipe == 0) return cpids;
 
   char buffer[256];
@@ -129,7 +130,7 @@ bool valid_monitor_disable(const std::string disable_name) {
     return false;
   }
   auto monitors = registry::Registry<Imonitor>::list_registered();
-  for (const auto& monitor_name : monitors) {
+  for (const auto &monitor_name : monitors) {
     if (monitor_name == disable_name) {
       return true;
     }
@@ -139,4 +140,31 @@ bool valid_monitor_disable(const std::string disable_name) {
   return false;
 }
 
+// Look for any monitors we should disable from the environment
+void disable_monitors_from_env(std::vector<std::string> &disabled_monitors) {
+  char *env_string = std::getenv("PRMON_DISABLE_MONITOR");
+  if (!env_string) return;
+
+  // We split this string on commas
+  unsigned pos{0}, start{0};
+  while (env_string[pos] != '\0') {
+    if (env_string[pos] == ',') {
+      snip_string_and_test(env_string, start, pos, disabled_monitors);
+      start = ++pos;
+    } else {
+      ++pos;
+    }
+  }
+  snip_string_and_test(env_string, start, pos, disabled_monitors);
+}
+
+// Snip a substring from the environment variable c-string
+// and test if it's a valid monitor name
+void snip_string_and_test(char *env_string, unsigned start, unsigned pos,
+                          std::vector<std::string> &disabled_monitors) {
+  std::string monitor_name(env_string + start, pos - start);
+  if (valid_monitor_disable(monitor_name))
+    disabled_monitors.push_back(monitor_name);
+}
+
 }  // namespace prmon

package/src/prmonutils.h

@@ -24,6 +24,11 @@ std::vector<pid_t> offspring_pids(const pid_t mother_pid);
 // Switch on/off states for individual monitors
 bool valid_monitor_disable(const std::string disable_name);
 
+// Disable monitors from an environment variable
+void disable_monitors_from_env(std::vector<std::string> &disabled_monitors);
+void snip_string_and_test(char *env_string, unsigned start, unsigned pos,
+                          std::vector<std::string> &disabled_monitors);
+
 // Signal handlers
 extern bool sigusr1;
 void SignalCallbackHandler(int);

package/tests/test_disable.py

@@ -7,6 +7,7 @@
 
 import argparse
 import json
+import os
 import subprocess
 import sys
 import unittest
@@ -18,13 +19,25 @@ def setup_configurable_test(disable=[]):
     class ConfigurableProcessMonitor(unittest.TestCase):
         """Test class for a specific set of parameters"""
 
-        def test_run_test_with_params(self):
-            """Actual test runner"""
+        __monitor_fields__ = {
+            "countmon": "nprocs",
+            "cpumon": "stime",
+            "iomon": "rchar",
+            "memmon": "vmem",
+            "netmon": "rx_bytes",
+            "nvidiamon": "ngpus",
+        }
+
+        def setup_and_run(self, envdisable):
+            """Baseline test run"""
             burn_cmd = ["./burner", "--time", "3"]
 
             prmon_cmd = ["../prmon", "--interval", "1"]
-            for d in disable:
-                prmon_cmd.extend(("-d", d))
+            if envdisable:
+                os.environ["PRMON_DISABLE_MONITOR"] = ",".join(disable)
+            else:
+                for d in disable:
+                    prmon_cmd.extend(("-d", d))
             prmon_cmd.append("--")
             prmon_cmd.extend(burn_cmd)
             prmon_p = subprocess.Popen(prmon_cmd, shell=False)
@@ -37,17 +50,20 @@ def test_run_test_with_params(self):
                 prmon_json = json.load(infile)
 
                 # Check we don't have fields corresponding to disabled monitors
-                monitor_fields = {
-                    "countmon": "nprocs",
-                    "cpumon": "stime",
-                    "iomon": "rchar",
-                    "memmon": "vmem",
-                    "netmon": "rx_bytes",
-                    "nvidiamon": "ngpus",
-                }
                 for d in disable:
-                    if d in monitor_fields:
-                        self.assertFalse(monitor_fields[d] in prmon_json["Max"])
+                    if d in ConfigurableProcessMonitor.__monitor_fields__:
+                        self.assertFalse(
+                            ConfigurableProcessMonitor.__monitor_fields__[d]
+                            in prmon_json["Max"]
+                        )
+
+        def test_disable_cli(self):
+            """Thin wrapper around selective test, CLI version"""
+            self.setup_and_run(False)
+
+        def test_disable_envvar(self):
+            """Thin wrapper around selective test, envvar version"""
+            self.setup_and_run(True)
 
     return ConfigurableProcessMonitor
 
@@ -56,7 +72,9 @@ def main_parse_args_and_get_test():
     """Parse arguments and call test class generator
     returning the test case (which is unusual for a
     main() function)"""
-    parser = argparse.ArgumentParser(description="Configurable test runner")
+    parser = argparse.ArgumentParser(
+        description="Configurable test runner - disable monitors"
+    )
     parser.add_argument("--disable", nargs="+")
 
     args = parser.parse_args()