Repository: terrapower/armi
Pull request #2018 (2024-11-21 21:05:55)

Improving calculation for total job memory
## What is the change?

This swaps the calculation for total job memory to use a more straightforward cluster resource request parameter. Some edits and additions to the tests flow from the update.

## Why is the change being made?

This makes the math simpler!

---

## Checklist

<!--
    You (the pull requester) should put an `x` in the boxes below you have completed.

    (If a checkbox doesn't apply to your PR, check it anyway.)
-->

- [x] This PR has only [one purpose or idea](https://terrapower.github.io/armi/developer/tooling.html#one-idea-one-pr).
- [x] [Tests](https://terrapower.github.io/armi/developer/tooling.html#test-it) have been added/updated to verify any new/changed code.

<!-- Check the code quality -->

- [x] The code style follows [good practices](https://terrapower.github.io/armi/developer/standards_and_practices.html).
- [x] The commit message(s) follow [good practices](https://terrapower.github.io/armi/developer/tooling.html).

<!-- Check the project-level cruft -->

- [x] ~The [release notes](https://terrapower.github.io/armi/developer/tooling.html#add-release-notes) have been updated if necessary.~ too small to matter
- [x] The [documentation](https://terrapower.github.io/armi/developer/tooling.html#document-it) is still up-to-date in the `doc` folder.
- [x] The dependencies are still up-to-date in `pyproject.toml`.

john-science: Wait, neither of these settings exist in ARMI:

```python
self.cs["nTasksPerNode"]  # OLD
self.cs["cpusPerTask"]  # NEW
```

If we are modifying this code, I would prefer to move it out of ARMI.
opotowsky: > Can you update the docstring for test_getTotalJobMemory?

I'm not sure how it's inaccurate? 
opotowsky: > Wait, neither of these settings exist in ARMI:
> 
> ```python
> self.cs["nTasksPerNode"]  # OLD
> self.cs["cpusPerTask"]  # NEW
> ```
> 
> If we are modifying this code, I would prefer to move it out of ARMI.

Discussed offline. Removal will happen! Just not in this PR.
albeanth: > > Can you update the docstring for test_getTotalJobMemory?
> 
> I'm not sure how it's inaccurate? It's the same test as before, just with the new inputs. (and I added a combo or two)

you're right! I was wrong. All good! 

Patch:
diff --git a/armi/bookkeeping/memoryProfiler.py b/armi/bookkeeping/memoryProfiler.py
index 12f9aebe..f3b6106d 100644
--- a/armi/bookkeeping/memoryProfiler.py
+++ b/armi/bookkeeping/memoryProfiler.py
@@ -36,7 +36,6 @@ See Also
 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 @@ class MemoryProfiler(interfaces.Interface):
     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/mpiActions.py b/armi/mpiActions.py
index 96b3495b..8026ddf1 100644
--- a/armi/mpiActions.py
+++ b/armi/mpiActions.py
@@ -505,7 +505,7 @@ class DistributeStateAction(MpiAction):
         self._skipInterfaces = skipInterfaces
 
     def invokeHook(self):
-        r"""Sync up all nodes with the reactor, the cs, and the interfaces.
+        """Sync up all nodes with the reactor, the cs, and the interfaces.
 
         Notes
         -----
@@ -604,9 +604,7 @@ class DistributeStateAction(MpiAction):
 
         self.r.o = self.o
 
-        runLog.debug(
-            "The reactor has {} assemblies".format(len(self.r.core.getAssemblies()))
-        )
+        runLog.debug(f"The reactor has {len(self.r.core.getAssemblies())} assemblies")
         # attach here so any interface actions use a properly-setup reactor.
         self.o.reattach(self.r, cs)  # sets r and cs
 
diff --git a/armi/physics/neutronics/macroXSGenerationInterface.py b/armi/physics/neutronics/macroXSGenerationInterface.py
index 18ecb33a..2665b82c 100644
--- a/armi/physics/neutronics/macroXSGenerationInterface.py
+++ b/armi/physics/neutronics/macroXSGenerationInterface.py
@@ -25,8 +25,8 @@ from armi import interfaces
 from armi import mpiActions
 from armi import runLog
 from armi.nuclearDataIO import xsCollections
-from armi.utils import iterables
 from armi.physics.neutronics.settings import CONF_MINIMUM_NUCLIDE_DENSITY
+from armi.utils import getBurnSteps, iterables
 
 
 class MacroXSGenerator(mpiActions.MpiAction):
@@ -112,12 +112,12 @@ class MacroXSGenerator(mpiActions.MpiAction):
 
 class MacroXSGenerationInterface(interfaces.Interface):
     """
-    Builds macroscopic cross sections on all blocks.
-
-    .. warning::
-         This probably shouldn't be an interface since it has no interactXYZ methods
-         It should probably be converted to an MpiAction.
+    Builds macroscopic cross sections on all Blocks.
 
+    Warning
+    -------
+    This probably shouldn't be an interface since it has no interactXYZ methods. It should probably
+    be converted to an MpiAction.
     """
 
     name = "macroXsGen"
@@ -136,8 +136,7 @@ class MacroXSGenerationInterface(interfaces.Interface):
         libType="micros",
     ):
         """
-        Builds block-level macroscopic cross sections for making diffusion
-        equation matrices.
+        Builds block-level macroscopic cross sections for making diffusion equation matrices.
 
         This will use MPI if armi.context.MPI_SIZE > 1
 
@@ -188,8 +187,9 @@ class MacroXSGenerationInterface(interfaces.Interface):
             block: either "micros" for neutron XS or "gammaXS" for gamma XS.
         """
         cycle = self.r.p.cycle
+        burnSteps = getBurnSteps(self.cs)
         self.macrosLastBuiltAt = (
-            sum([self.r.o.burnSteps[i] + 1 for i in range(cycle)]) + self.r.p.timeNode
+            sum([burnSteps[i] + 1 for i in range(cycle)]) + self.r.p.timeNode
         )
 
         runLog.important("Building macro XS")
@@ -205,10 +205,8 @@ class MacroXSGenerationInterface(interfaces.Interface):
         xsGen.invoke(self.o, self.r, self.cs)
 
 
-# helper functions for mpi communication
-
-
 def _scatterList(lst):
+    """Helper functions for mpi communication."""
     if context.MPI_RANK == 0:
         chunked = iterables.split(lst, context.MPI_SIZE)
     else:
@@ -217,6 +215,7 @@ def _scatterList(lst):
 
 
 def _gatherList(localList):
+    """Helper functions for mpi communication."""
     globalList = context.MPI_COMM.gather(localList, root=0)
     if context.MPI_RANK == 0:
         globalList = iterables.flatten(globalList)
diff --git a/armi/reactor/assemblies.py b/armi/reactor/assemblies.py
index 73503efc..29e8d8f9 100644
--- a/armi/reactor/assemblies.py
+++ b/armi/reactor/assemblies.py
@@ -1018,28 +1018,24 @@ class Assembly(composites.Composite):
         This caches interpolators for each param and must be cleared if new params are
         set or new heights are set.
 
-        WARNING:
-        Fails when requested to extrapolate.With higher order splines it is possible
-        to interpolate non-physical values, for example a negative flux or dpa. Please
-        use caution when going off default in interpType and be certain that
-        interpolated values are physical.
+        Warning
+        -------
+        Fails when requested to extrapolate. With higher order splines it is possible to interpolate
+        non-physical values, for example, a negative flux or dpa. Please use caution when going off
+        default in interpType and be certain that interpolated values are physical.
 
         Parameters
         ----------
         param : str
             the parameter to interpolate
-
         elevations : array of float
-            the elevations from the bottom of the assembly in cm at which you want the
-            point.
-
+            the elevations from the bottom of the assembly in cm at which you want the point.
         interpType: str or int
             used in interp1d. interp1d documention: Specifies the kind of interpolation
             as a string ('linear', 'nearest', 'zero', 'slinear', 'quadratic', 'cubic'
             where 'slinear', 'quadratic' and 'cubic' refer to a spline interpolation of
             first, second or third order) or as an integer specifying the order of the
             spline interpolator to use. Default is 'linear'.
-
         fillValue: str
             Rough pass through to scipy.interpolate.interp1d. If 'extend', then the
             lower and upper bounds are used as the extended value. If 'extrapolate',
@@ -1070,23 +1066,22 @@ class Assembly(composites.Composite):
         This caches interpolators for each param and must be cleared if new params are
         set or new heights are set.
 
-        WARNING: Fails when requested to extrapololate. With higher order splines it is
-        possible to interpolate nonphysical values, for example a negative flux or dpa.
-        Please use caution when going off default in interpType and be certain that
-        interpolated values are physical.
+        Warning
+        -------
+        Fails when requested to extrapololate. With higher order splines it is possible to
+        interpolate nonphysical values, for example, a negative flux or dpa. Please use caution when
+        going off default in interpType and be certain that interpolated values are physical.
 
         Parameters
         ----------
         param : str
             the parameter to interpolate
-
         interpType: str or int
             used in interp1d. interp1d documention: Specifies the kind of interpolation
             as a string ('linear', 'nearest', 'zero', 'slinear', 'quadratic', 'cubic'
             where 'slinear', 'quadratic' and 'cubic' refer to a spline interpolation of
             first, second or third order) or as an integer specifying the order of the
             spline interpolator to use. Default is 'linear'.
-
         fillValue: float
             Rough pass through to scipy.interpolate.interp1d. If 'extend', then the
             lower and upper bounds are used as the extended value. If 'extrapolate',
diff --git a/armi/reactor/converters/parameterSweeps/generalParameterSweepConverters.py b/armi/reactor/converters/parameterSweeps/generalParameterSweepConverters.py
index fa8308fa..11248e3b 100644
--- a/armi/reactor/converters/parameterSweeps/generalParameterSweepConverters.py
+++ b/armi/reactor/converters/parameterSweeps/generalParameterSweepConverters.py
@@ -52,15 +52,6 @@ class SettingsModifier(ParameterSweepConverter):
             )
 
 
-class CustomModifier(ParameterSweepConverter):
-    """Invoke the shuffle logic `applyCustomPerturbation` method to make a custom setting."""
-
-    def convert(self, r=None):
-        ParameterSweepConverter.convert(self, r)
-        fh = r.o.getInterface("fuelHandler")
-        fh.applyCustomPerturbation(self._parameter)
-
-
 class NeutronicConvergenceModifier(ParameterSweepConverter):
     """Adjusts the neutronics convergence parameters."""
 


Test patch:
diff --git a/armi/bookkeeping/tests/test_memoryProfiler.py b/armi/bookkeeping/tests/test_memoryProfiler.py
index 94854cf8..497fb3a6 100644
--- a/armi/bookkeeping/tests/test_memoryProfiler.py
+++ b/armi/bookkeeping/tests/test_memoryProfiler.py
@@ -139,9 +139,18 @@ class TestMemoryProfiler(unittest.TestCase):
         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 @@ class TestMemoryProfiler(unittest.TestCase):
         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()
diff --git a/armi/physics/fuelCycle/tests/test_fuelHandlers.py b/armi/physics/fuelCycle/tests/test_fuelHandlers.py
index c4aca76a..8fd0f1f4 100644
--- a/armi/physics/fuelCycle/tests/test_fuelHandlers.py
+++ b/armi/physics/fuelCycle/tests/test_fuelHandlers.py
@@ -481,7 +481,7 @@ class TestFuelHandler(FuelHandlerTestHelper):
             self.assertEqual(a.getLocation(), "SFP")
 
         # do some shuffles
-        fh = self.r.o.getInterface("fuelHandler")
+        fh = self.o.getInterface("fuelHandler")
         self.runShuffling(fh)  # changes caseTitle
 
         # Make sure the generated shuffles file matches the tracked one.  This will need to be
@@ -504,7 +504,7 @@ class TestFuelHandler(FuelHandlerTestHelper):
         newSettings["explicitRepeatShuffles"] = "armiRun-SHUFFLES.txt"
         self.o.cs = self.o.cs.modified(newSettings=newSettings)
 
-        fh = self.r.o.getInterface("fuelHandler")
+        fh = self.o.getInterface("fuelHandler")
 
         self.runShuffling(fh)
 
diff --git a/armi/reactor/converters/parameterSweeps/tests/test_paramSweepConverters.py b/armi/reactor/converters/parameterSweeps/tests/test_paramSweepConverters.py
index 4592cca6..97fc3d95 100644
--- a/armi/reactor/converters/parameterSweeps/tests/test_paramSweepConverters.py
+++ b/armi/reactor/converters/parameterSweeps/tests/test_paramSweepConverters.py
@@ -18,7 +18,6 @@ import unittest
 
 from armi.tests import TEST_ROOT
 from armi.reactor.converters.parameterSweeps.generalParameterSweepConverters import (
-    CustomModifier,
     NeutronicConvergenceModifier,
     ParameterSweepConverter,
     SettingsModifier,
@@ -64,15 +63,3 @@ class TestParamSweepConverters(unittest.TestCase):
         # NOTE: Settings objects are not modified, but we point to new objects
         self.assertIn("Simple test input", self.cs["comment"])
         self.assertEqual(con._cs["comment"], "FakeParam")
-
-    def test_customModifier(self):
-        """Super basic test of the Custom Modifier."""
-        con = CustomModifier(self.cs, "FakeParam")
-        self.assertEqual(con._parameter, "FakeParam")
-
-        # For testing purposes, we need to add a dummy perturbation
-        fh = self.r.o.getInterface("fuelHandler")
-        fh.applyCustomPerturbation = lambda val: val
-
-        con.convert(self.r)
-        self.assertEqual(con._sourceReactor, self.r)