Set memory for ignored tasks (Express Merge)

src/python/WMCore/WMSpec/WMTask.py

@@ -596,15 +596,19 @@ def jobSplittingParameters(self, performance=True):
 
         return splittingParams
 
-    def setJobResourceInformation(self, timePerEvent=None, sizePerEvent=None, memoryReq=None):
+    def setJobResourceInformation(self, timePerEvent=None, sizePerEvent=None, memoryReq=None, taskType=None):
         """
         _setJobResourceInformation_
 
         Set the values to estimate the required computing resources for a job,
         the three key values are main memory usage, time per processing unit (e.g. time per event) and
         disk usage per processing unit (e.g. size per event).
         """
-        if self.taskType() in ["Merge", "Cleanup", "LogCollect"]:
+
+        # If task type is specified, it may be that we dont want "Merge" tasks to be ignored, 
+        # or that the task type is not one of the three to ignore
+
+        if self.taskType() in ["Merge", "Cleanup", "LogCollect"] and taskType is None:
             # don't touch job requirements for these task types
             return
 
@@ -621,7 +625,7 @@ def setJobResourceInformation(self, timePerEvent=None, sizePerEvent=None, memory
         if memoryReq or getattr(performanceParams, "memoryRequirement", None):
             performanceParams.memoryRequirement = memoryReq or getattr(performanceParams, "memoryRequirement")
             # if we change memory requirements, then we must change MaxPSS as well
-            self.setMaxPSS(performanceParams.memoryRequirement)
+            self.setMaxPSS(performanceParams.memoryRequirement, taskType)
 
         return
 
@@ -1215,14 +1219,17 @@ def _setPerformanceMonitorConfig(self):
             self.monitoring.section_("PerformanceMonitor")
         return
 
-    def setMaxPSS(self, maxPSS):
+    def setMaxPSS(self, maxPSS, taskType=None):
         """
         _setMaxPSS_
 
         Set MaxPSS performance monitoring for this task.
         :param maxPSS: maximum Proportional Set Size (PSS) memory consumption in MiB
         """
-        if self.taskType() in ["Merge", "Cleanup", "LogCollect"]:
+        # If task type is specified, it may be that we dont want "Merge" tasks to be ignored, 
+        # or that the task type is not one of the three to ignore
+
+        if self.taskType() in ["Merge", "Cleanup", "LogCollect"] and taskType is None:
             # keep the default settings (from StdBase) for these task types
             return
 

src/python/WMCore/WMSpec/WMWorkload.py

@@ -896,7 +896,14 @@ def setMemory(self, memory, initialTask=None):
                 mem = memory.get(task.name())
             else:
                 mem = memory
-            task.setJobResourceInformation(memoryReq=mem)
+
+            # If a task is specified, the ResourceInformation should take into account that
+            # it may not want to be ignored, regardless of the task 
+            if initialTask is not None:
+                task.setJobResourceInformation(memoryReq=mem, taskType=task.taskType()) # Or taskType=initialTask.taskType() ?
+            else:
+                task.setJobResourceInformation(memoryReq=mem)
+
             self.setMemory(memory, task)
 
         return
@@ -2002,6 +2009,10 @@ def updateArguments(self, kwargs):
 
         if kwargs.get("Memory") is not None:
             self.setMemory(kwargs.get("Memory"))
+        if kwargs.get("TaskMemory") is not None:
+            for taskName in kwargs.get("TaskMemory"):
+                task = self.getTaskByName(taskName)
+                self.setMemory(memory=kwargs.get("TaskMemory"), initialTask=task)
         if kwargs.get("Multicore") is not None:
             self.setCoresAndStreams(kwargs.get("Multicore"), kwargs.get("EventStreams"))