UpdateSandbox modifications with new logging

dmwm · Mar 29, 2024 · 8b852b3 · 8b852b3
1 parent 1c6b822
commit 8b852b3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 9 deletions.
diff --git a/src/python/WMComponent/RetryManager/Modifier/BaseModifier.py b/src/python/WMComponent/RetryManager/Modifier/BaseModifier.py
@@ -21,8 +21,7 @@ class BaseModifier(object):
 
     def __init__(self, config):
         object.__init__(self)
-        self.backupPath = "./oldSandboxes"
-        #self.tempFolder --> not being used? Initialized as what?
+        self.backupPath = "oldSandboxes/"
         self.sandboxPath = None
         self.config = config
 
@@ -42,16 +41,27 @@ def loadJobPKL(self, pklFile):
 
     def updateSandbox(self, jobPKL, workload): # Not using workload?
         date = datetime.datetime.now().strftime("%y%m%d%H%M%S")
+        os.makedirs(os.path.dirname(self.backupPath), exist_ok=True)
         backupFile = f"{self.backupPath}/{jobPKL['workflow']}_{date}.tar.bz2"
 
         shutil.copyfile(jobPKL['sandbox'], backupFile)
 
         tempDir = TemporaryDirectory()
+        tempDirName = tempDir.name
+
         tFile = tarfile.open(jobPKL['sandbox'], "r")
-        tFile.extractall(tempDir)
+        tFile.extractall(tempDirName)
 
-        shutil.copyfile(jobPKL['spec'], tempDir+'/WMSandbox/WMWorkload.pkl')
+        shutil.copyfile(jobPKL['spec'], tempDirName+'/WMSandbox/WMWorkload.pkl')
+
+        archivePath = jobPKL['sandbox']
+        with tarfile.open(archivePath, "w:bz2") as tar:
+            for folder in os.listdir(tempDirName):
+                tar.add(f"{tempDirName}/{folder}", arcname=folder)
 
+        tempDir.cleanup()
+        return
+
 
     def getWorkload(self, jobPKL):
         """

diff --git a/src/python/WMComponent/RetryManager/Modifier/MemoryModifier.py b/src/python/WMComponent/RetryManager/Modifier/MemoryModifier.py
@@ -50,14 +50,14 @@ def checkNewJobPkl(pklFile):
         print (data['estimatedMemoryUsage'])
 
     def getNewMemory(self, jobPKL, settings):
-        maxMemPerCore = settings['maxMemory']/jobPKL['numberOfCores']
+        maxMemPerCore = settings['maxMemory']
         currentMem = jobPKL['estimatedMemoryUsage']
         currentMemPerCore = currentMem/jobPKL['numberOfCores']
 
         if 'multiplyMemory' in settings:
             newMemPerCore = currentMemPerCore * settings['multiplyMemory']
         elif 'addMemory' in settings:
-            newMemPerCore = currentMemPerCore + (settings['addMemory']/jobPKL['numberOfCores'])
+            newMemPerCore = currentMemPerCore + settings['addMemory']
         else:
             newMemPerCore = currentMemPerCore
             logging.info('No increment values were given in the MemoryModifier parameter')
@@ -80,14 +80,15 @@ def changeMemory(self, job, settings):
 
         self.changeJobPkl(pklFile, jobPKL, newMemory)
         self.changeSandbox(jobPKL, newMemory)
+        logging.info('Old maxPSS: %d. New maxPSS: %d', job['estimatedMemoryUsage'], newMemory)
 
     def modifyJob(self, job):
         try:
             settings = self.getModifierParam(job['jobType'], 'settings')
         except:
             logging.exception('Error while getting the MemoryModifier settings parameter. Not modifying memory')
             return
-        
+
         if not 'requiresModify' in settings:
             logging.info('requiresModify not specified')
             logging.info('Not performing any modifications')
@@ -99,5 +100,5 @@ def modifyJob(self, job):
             return
 
         else:
-            logging.info('MemoryModifier.changeMemory called successfully')
-            self.changeMemory(job, settings)
+            logging.info('Modifying memory for job %d', job['id'])
+            self.changeMemory(job, settings)