Enable getting user usage of CPU microsecond from cgroup v2

Author: appolloford

sams/sampler/SlurmCGroup.py

@@ -152,15 +152,29 @@ def _get_cgroup_item_path(self, resource_type, value):
         return os.path.join(self.cgroup_base, resource_type,
                             self.cgroup, value)
 
-    def read_cgroup(self, *items):
+    def read_cgroup(self, *items, metric_name=None):
+        """Get metric value from target file. If `metric_name` is provided,
+        file will be parsed to return the first found value. Otherwise the file
+        should have one line indicating the value.
+        """
         path = self._get_cgroup_item_path(*items)
         try:
             with open(path, "r") as file:
-                return file.readline().strip()
+                if not metric_name:
+                    return file.readline().strip()
+                else:
+                    for line in file:
+                        if line.startswith(metric_name):
+                            return line.split()[1].strip()
+                    else:
+                        raise RuntimeError(f"Failed to extract {metric_name} value from {path}")
         except OSError as err:
             logger.error(f"Failed to open {path} for reading")
             logger.error(err)
             return ""
+        except RuntimeError as err:
+            logger.error(err)
+            return ""
 
     @classmethod
     def final_data(cls):

sams/sampler/SlurmCGroup2.py

@@ -32,6 +32,7 @@ def sample(self):
         logger.debug("sample()")
 
         cpus = self._cpucount(self.read_cgroup("cpuset.cpus"))
+        user_cpu_usec = self.read_cgroup("cpu.stat", metric_name="user_usec")
         memory_usage = self.read_cgroup("memory.current")
         memory_limit = self.read_cgroup("memory.high")
         memory_max_usage = self.read_cgroup("memory.max")
@@ -40,10 +41,12 @@ def sample(self):
 
         entry = {
             "cpus": cpus,
+            "user_cpu_usec": user_cpu_usec,
             "memory_usage": memory_usage,
             "memory_limit": memory_limit,
             "memory_max_usage": memory_max_usage,
-            "memory_swap": str(int(memory_usage_and_swap) - int(memory_usage))}
+            "memory_swap": str(int(memory_usage_and_swap) - int(memory_usage)),
+        }
         self.compute_sample_averages(entry)
         self._most_recent_sample = [self._storage_wrapping(entry)]
         self.store(entry)