Improve the _GetTotalMemoryKbFromCgroup function to read from

PiperOrigin-RevId: 568596288
GoogleCloudPlatform · Sep 26, 2023 · dd1d4a7 · dd1d4a7
1 parent ab4296c
commit dd1d4a7
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 5 deletions.
diff --git a/perfkitbenchmarker/linux_virtual_machine.py b/perfkitbenchmarker/linux_virtual_machine.py
@@ -1540,7 +1540,8 @@ def _GetTotalMemoryKbFromCgroup(self):
     """Extracts the memory space in kibibyte (KiB) for containers.
 
     Gets the memory capacity from
-    /sys/fs/cgroup/memory/<container>/memory.limit_in_bytes.
+    /sys/fs/cgroup/memory/<container>/memory.limit_in_bytes,
+    or /sys/fs/cgroup/memory/memory.limit_in_bytes.
     Below are the example of their returns:
     $ cat /sys/fs/cgroup/memory/container/memory.limit_in_bytes
       1024
@@ -1549,8 +1550,9 @@ def _GetTotalMemoryKbFromCgroup(self):
       The memory capacity in kibibyte (KiB).
 
     Raises:
-      ValueError: If not found /proc/self/cgroup, or
-      /sys/fs/cgroup/memory/<container>/memory.limit_in_bytes.
+      ValueError: If not found /proc/self/cgroup,
+      or /sys/fs/cgroup/memory/<container>/memory.limit_in_bytes,
+      or /sys/fs/cgroup/memory/memory.limit_in_bytes.
     """
     if self._RemoteFileExists('/proc/self/cgroup'):
       container_name, _ = self.RemoteCommand(
@@ -1569,10 +1571,16 @@ def _GetTotalMemoryKbFromCgroup(self):
           f'cat /sys/fs/cgroup/memory/{container_name}/memory.limit_in_bytes'
       )
       return int(stdout) // 1024
+    elif self._RemoteFileExists('/sys/fs/cgroup/memory/memory.limit_in_bytes'):
+      stdout, _ = self.RemoteCommand(
+          'cat /sys/fs/cgroup/memory/memory.limit_in_bytes'
+      )
+      return int(stdout) // 1024
 
     raise ValueError(
-        '_GetTotalMemoryKbFromCgroup failed, '
-        'cannot read /sys/fs/cgroup/memory/<container>/memory.limit_in_bytes.'
+        '_GetTotalMemoryKbFromCgroup failed, cannot read '
+        ' /sys/fs/cgroup/memory/%s/memory.limit_in_bytes or'
+        ' /sys/fs/cgroup/memory/memory.limit_in_bytes' % container_name
     )
 
   def _GetTotalMemoryKb(self):

diff --git a/tests/linux_virtual_machine_test.py b/tests/linux_virtual_machine_test.py
@@ -642,6 +642,22 @@ def testMemCapFromDifferentSources(self):
     }
     vm = self.CreateVm(responses)
     self.assertEqual(1, vm.total_memory_kb)
+    responses = {
+        'ls /proc/self/cgroup >> /dev/null 2>&1 || echo file_not_exist': '',
+        "grep memory /proc/self/cgroup |cut -d ':' -f 3 |sed -e 's:^/::'": (
+            'container\n'
+        ),
+        'ls /sys/fs/cgroup/memory/container/memory.limit_in_bytes >> /dev/null 2>&1 || echo file_not_exist': (
+            'file_not_exist'
+        ),
+        'ls /sys/fs/cgroup/memory/memory.limit_in_bytes >> /dev/null 2>&1 || echo file_not_exist': (
+            ''
+        ),
+        'cat /sys/fs/cgroup/memory/memory.limit_in_bytes': '2048',
+        "cat /proc/meminfo | grep MemTotal | awk '{print $2}'": '3',
+    }
+    vm = self.CreateVm(responses)
+    self.assertEqual(2, vm.total_memory_kb)
     FLAGS['use_cgroup_memory_limits'].parse(False)
     responses = {
         "cat /proc/meminfo | grep MemTotal | awk '{print $2}'": '3',