diff --git a/armi/bookkeeping/memoryProfiler.py b/armi/bookkeeping/memoryProfiler.py index 12f9aebef..f3b6106d4 100644 --- a/armi/bookkeeping/memoryProfiler.py +++ b/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( diff --git a/armi/bookkeeping/tests/test_memoryProfiler.py b/armi/bookkeeping/tests/test_memoryProfiler.py index 94854cf85..497fb3a64 100644 --- a/armi/bookkeeping/tests/test_memoryProfiler.py +++ b/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()