Improving calculation for total job memory (#2018)

Author: opotowsky

armi/bookkeeping/memoryProfiler.py

@@ -36,7 +36,6 @@
 https://pythonhosted.org/psutil/
 https://docs.python.org/3/library/gc.html#gc.garbage
 """
-from math import floor
 from os import cpu_count
 from typing import Optional
 import gc
@@ -71,14 +70,11 @@ def describeInterfaces(cs):
     return (MemoryProfiler, {})
 
 
-def getTotalJobMemory(nTasksPerNode):
+def getTotalJobMemory(nTasks, cpusPerTask):
     """Function to calculate the total memory of a job. This is a constant during a simulation."""
     cpuPerNode = cpu_count()
     ramPerCpuGB = psutil.virtual_memory().total / (1024**3) / cpuPerNode
-    if nTasksPerNode == 0:
-        nTasksPerNode = cpuPerNode
-    cpusPerTask = floor(cpuPerNode / nTasksPerNode)
-    jobMem = nTasksPerNode * cpusPerTask * ramPerCpuGB
+    jobMem = nTasks * cpusPerTask * ramPerCpuGB
     return jobMem
 
 
@@ -137,14 +133,15 @@ def interactEOL(self):
     def printCurrentMemoryState(self):
         """Print the current memory footprint and available memory."""
         try:
-            nTasksPerNode = self.cs["nTasksPerNode"]
+            cpusPerTask = self.cs["cpusPerTask"]
         except NonexistentSetting:
             runLog.extra(
                 "To view memory consumed, remaining available, and total allocated for a case, "
-                "add the setting 'nTasksPerNode' to your application."
+                "add the setting 'cpusPerTask' to your application."
             )
             return
-        totalMemoryInGB = getTotalJobMemory(nTasksPerNode)
+        nTasks = self.cs["nTasks"]
+        totalMemoryInGB = getTotalJobMemory(nTasks, cpusPerTask)
         currentMemoryUsageInGB = getCurrentMemoryUsage() / 1024
         availableMemoryInGB = totalMemoryInGB - currentMemoryUsageInGB
         runLog.info(

armi/bookkeeping/tests/test_memoryProfiler.py

@@ -139,9 +139,18 @@ def test_getTotalJobMemory(self, mockCpuCount, mockVMem):
         vMem.total = (1024**3) * 50
         mockVMem.return_value = vMem
 
-        expectedArrangement = {0: 50, 1: 50, 2: 50, 3: 45, 4: 40, 5: 50}
-        for nTasksPerNode, jobMemory in expectedArrangement.items():
-            self.assertEqual(getTotalJobMemory(nTasksPerNode), jobMemory)
+        expectedArrangement = {
+            (10, 1): 50,
+            (1, 10): 50,
+            (2, 5): 50,
+            (3, 3): 45,
+            (4, 1): 20,
+            (2, 4): 40,
+            (5, 2): 50,
+        }
+        for compReq, jobMemory in expectedArrangement.items():
+            # compReq[0] is nTasks and compReq[1] is cpusPerTask
+            self.assertEqual(getTotalJobMemory(compReq[0], compReq[1]), jobMemory)
 
     @patch("armi.bookkeeping.memoryProfiler.PrintSystemMemoryUsageAction")
     @patch("armi.bookkeeping.memoryProfiler.SystemAndProcessMemoryUsage")
@@ -166,20 +175,26 @@ def test_printCurrentMemoryState(
         mockVMem.return_value = vMem
         self._setMemUseMock(mockPrintSysMemUseAction)
         with mockRunLogs.BufferLog() as mockLogs:
-            csMock = MagicMock()
-            csMock.__getitem__.return_value = 2
-            self.memPro.cs = csMock
+            self.memPro.cs = {"cpusPerTask": 1, "nTasks": 10}
             self.memPro.printCurrentMemoryState()
             stdOut = mockLogs.getStdout()
             self.assertIn("Currently using 6.0 GB of memory.", stdOut)
             self.assertIn("There is 44.0 GB of memory left.", stdOut)
             self.assertIn("There is a total allocation of 50.0 GB", stdOut)
+            # Try another for funzies where we only use half the available resources on the node
+            mockLogs.emptyStdout()
+            self.memPro.cs = {"cpusPerTask": 5, "nTasks": 1}
+            self.memPro.printCurrentMemoryState()
+            stdOut = mockLogs.getStdout()
+            self.assertIn("Currently using 6.0 GB of memory.", stdOut)
+            self.assertIn("There is 19.0 GB of memory left.", stdOut)
+            self.assertIn("There is a total allocation of 25.0 GB", stdOut)
 
     def test_printCurrentMemoryState_noSetting(self):
         """Test that the try/except works as it should."""
         expectedStr = (
             "To view memory consumed, remaining available, and total allocated for a case, "
-            "add the setting 'nTasksPerNode' to your application."
+            "add the setting 'cpusPerTask' to your application."
         )
         with mockRunLogs.BufferLog() as mockLogs:
             self.memPro.printCurrentMemoryState()