Repository: DataBiosphere/toil

Issue:
Add option for retaining files in work directory (--workDir)
Allow users a way to retain temporary files generated during a toil run. Should probably be a new option (not --clean) since it will be very detrimental to toil's performance and only make sense in a debugging situation.

Add option for retaining files in work directory (--workDir)
Allow users a way to retain temporary files generated during a toil run. Should probably be a new option (not --clean) since it will be very detrimental to toil's performance and only make sense in a debugging situation.



Discussion:

Resolving patch:
diff --git a/src/toil/common.py b/src/toil/common.py
--- a/src/toil/common.py
+++ b/src/toil/common.py
@@ -44,6 +44,7 @@ def __init__(self):
         # Because the stats option needs the jobStore to persist past the end of the run,
         # the clean default value depends the specified stats option and is determined in setOptions
         self.clean = None
+        self.cleanWorkDir = None
 
         #Restarting the workflow options
         self.restart = False
@@ -123,6 +124,7 @@ def fC(minValue, maxValue=None):
         #TODO: LOG LEVEL STRING
         setOption("workDir")
         setOption("stats")
+        setOption("cleanWorkDir")
         setOption("clean")
         if self.stats:
             if self.clean != "never" and self.clean is not None:
@@ -197,7 +199,11 @@ def _addOptions(addGroupFn, config):
                             "information from the jobStore upon completion so the jobStore will never be deleted with"
                             "that flag. If you wish to be able to restart the run, choose \'never\' or \'onSuccess\'. "
                             "Default is \'never\' if stats is enabled, and \'onSuccess\' otherwise"))
-
+    addOptionFn("--cleanWorkDir", dest="cleanWorkDir",
+                choices=['always', 'never', 'onSuccess', 'onError'], default='always',
+                help=("Determines deletion of temporary worker directory upon completion of a job. Choices: 'always', "
+                      "'never', 'onSuccess'. Default = always. WARNING: This option should be changed for debugging "
+                      "only. Running a full pipeline with this option could fill your disk with intermediate data."))
     #
     #Restarting the workflow options
     #
diff --git a/src/toil/leader.py b/src/toil/leader.py
--- a/src/toil/leader.py
+++ b/src/toil/leader.py
@@ -521,7 +521,11 @@ def mainLoop(config, batchSystem, jobStore, rootJobWrapper, jobCache=None):
     with jobStore.readSharedFileStream("rootJobReturnValue") as fH:
         jobStoreFileID = fH.read()
     with jobStore.readFileStream(jobStoreFileID) as fH:
-        rootJobReturnValue = cPickle.load(fH)
+        try:
+            rootJobReturnValue = cPickle.load(fH)
+        except EOFError:
+            logger.exception("Failed to unpickle root job return value")
+            raise FailedJobsException(jobStoreFileID, totalFailedJobs)
     
     if totalFailedJobs > 0:
         if config.clean == "onError" or config.clean == "always" :
diff --git a/src/toil/worker.py b/src/toil/worker.py
--- a/src/toil/worker.py
+++ b/src/toil/worker.py
@@ -32,20 +32,10 @@
 
 logger = logging.getLogger( __name__ )
 
-def truncateFile(fileNameString, tooBig=50000):
-    """
-    Truncates a file that is bigger than tooBig bytes, leaving only the 
-    last tooBig bytes in the file.
-    """
-    if os.path.getsize(fileNameString) > tooBig:
-        fh = open(fileNameString, 'rb+')
-        fh.seek(-tooBig, 2) 
-        data = fh.read()
-        fh.seek(0) # rewind
-        fh.write(data)
-        fh.truncate()
-        fh.close()
-    
+
+logFileByteReportLimit = 50000
+
+
 def nextOpenDescriptor():
     """Gets the number of the next available file descriptor.
     """
@@ -246,7 +236,7 @@ def badWorker():
             
         #Make a temporary file directory for the jobWrapper
         localTempDir = makePublicDir(os.path.join(localWorkerTempDir, "localTempDir"))
-    
+
         ##########################################
         #Setup the stats, if requested
         ##########################################
@@ -393,8 +383,11 @@ def badWorker():
             statsDict.workers.clock = str(totalCPUTime - startClock)
             statsDict.workers.memory = str(totalMemoryUsage)
 
-        statsDict.workers.logsToMaster = messages  # logToMaster messages should be always be passed
+		# logToMaster messages should be always be passed
+        statsDict.workers.logsToMaster = messages
 
+        # log the worker log path here so that if the file is truncated the path can still be found
+        logger.info("Worker log can be found at %s. Set --cleanWorkDir to retain this log", localTempDir)
         logger.info("Finished running the chain of jobs on this node, we ran for a total of %f seconds", time.time() - startTime)
     
     ##########################################
@@ -451,21 +444,28 @@ def badWorker():
     
     #Copy back the log file to the global dir, if needed
     if workerFailed:
-        truncateFile(tempWorkerLogPath)
-        jobWrapper.logJobStoreFileID = jobStore.writeFile( tempWorkerLogPath, jobWrapper.jobStoreID )
-        os.remove(tempWorkerLogPath)
+        jobWrapper.logJobStoreFileID = jobStore.getEmptyFileStoreID(jobWrapper.jobStoreID)
+        with jobStore.updateFileStream(jobWrapper.logJobStoreFileID) as w:
+            with open(tempWorkerLogPath, "r") as f:
+                if os.path.getsize(tempWorkerLogPath) > logFileByteReportLimit:
+                    f.seek(-logFileByteReportLimit, 2)  # seek to last tooBig bytes of file
+                w.write(f.read())
         jobStore.update(jobWrapper)
-    elif debugging: # write log messages
-        truncateFile(tempWorkerLogPath)
+
+    elif debugging:  # write log messages
         with open(tempWorkerLogPath, 'r') as logFile:
+            if os.path.getsize(tempWorkerLogPath) > logFileByteReportLimit:
+                logFile.seek(-logFileByteReportLimit, 2)  # seek to last tooBig bytes of file
             logMessages = logFile.read().splitlines()
-        statsDict.logs = [Expando(jobStoreID=jobStoreID,text=logMessage) for logMessage in logMessages]
+        statsDict.logs = [Expando(jobStoreID=jobStoreID, text=logMessage) for logMessage in logMessages]
 
-    if (debugging or config.stats or messages) and not workerFailed: # We have stats/logging to report back
+    if (debugging or config.stats or messages) and not workerFailed:  # We have stats/logging to report back
         jobStore.writeStatsAndLogging(json.dumps(statsDict))
 
     #Remove the temp dir
-    shutil.rmtree(localWorkerTempDir)
+    cleanUp = config.cleanWorkDir
+    if cleanUp == 'always' or (cleanUp == 'onSuccess' and not workerFailed) or (cleanUp == 'onError' and workerFailed):
+        shutil.rmtree(localWorkerTempDir)
     
     #This must happen after the log file is done with, else there is no place to put the log
     if (not workerFailed) and jobWrapper.command == None and len(jobWrapper.stack) == 0:


Test patch:
diff --git a/src/toil/test/src/retainTempDirTest.py b/src/toil/test/src/retainTempDirTest.py
new file mode 100644
--- /dev/null
+++ b/src/toil/test/src/retainTempDirTest.py
@@ -0,0 +1,98 @@
+# Copyright (C) 2015 UCSC Computational Genomics Lab
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import absolute_import, print_function
+import os
+import shutil
+from toil.job import Job
+from toil.leader import FailedJobsException
+from toil.test import ToilTest
+
+class CleanWorkDirTest(ToilTest):
+    """
+    Tests testing the Job.FileStore class
+    """
+    def setUp(self):
+        super(CleanWorkDirTest, self).setUp()
+        self.testDir = self._createTempDir()
+
+    def tearDown(self):
+        super(CleanWorkDirTest, self).tearDown()
+        shutil.rmtree(self.testDir)
+
+    def testNever(self):
+        retainedTempData = self._runAndReturnWorkDir("never", job=tempFileTestJob)
+        self.assertNotEqual(retainedTempData, [], "The worker's temporary workspace was deleted despite "
+                                                  "cleanWorkDir being set to 'never'")
+
+    def testAlways(self):
+        retainedTempData = self._runAndReturnWorkDir("always", job=tempFileTestJob)
+        self.assertEqual(retainedTempData, [], "The worker's temporary workspace was not deleted despite "
+                                               "cleanWorkDir being set to 'always'")
+
+    def testOnErrorWithError(self):
+        retainedTempData = self._runAndReturnWorkDir("onError", job=tempFileTestErrorJob, expectError=True)
+        self.assertEqual(retainedTempData, [], "The worker's temporary workspace was not deleted despite "
+                                               "an error occurring and cleanWorkDir being set to 'onError'")
+
+    def testOnErrorWithNoError(self):
+        retainedTempData = self._runAndReturnWorkDir("onError", job=tempFileTestJob)
+        self.assertNotEqual(retainedTempData, [], "The worker's temporary workspace was deleted despite "
+                                                  "no error occurring and cleanWorkDir being set to 'onError'")
+
+    def testOnSuccessWithError(self):
+        retainedTempData = self._runAndReturnWorkDir("onSuccess", job=tempFileTestErrorJob, expectError=True)
+        self.assertNotEqual(retainedTempData, [], "The worker's temporary workspace was deleted despite "
+                                                  "an error occurring and cleanWorkDir being set to 'onSuccesss'")
+
+    def testOnSuccessWithSuccess(self):
+        retainedTempData = self._runAndReturnWorkDir("onSuccess", job=tempFileTestJob)
+        self.assertEqual(retainedTempData, [], "The worker's temporary workspace was not deleted despite "
+                                               "a successful job execution and cleanWorkDir being set to 'onSuccesss'")
+
+    def _runAndReturnWorkDir(self, cleanWorkDir, job, expectError=False):
+        """
+        Runs toil with the specified job and cleanWorkDir setting. expectError determines whether the test's toil
+        run is expected to succeed, and the test will fail if that expectation is not met. returns the contents of
+        the workDir after completion of the run
+        """
+        options = Job.Runner.getDefaultOptions(self._getTestJobStorePath())
+        options.workDir = self.testDir
+        options.clean = "always"
+        options.cleanWorkDir = cleanWorkDir
+        A = Job.wrapJobFn(job)
+        if expectError:
+            self._launchError(A, options)
+        else:
+            self._launchRegular(A, options)
+        return os.listdir(self.testDir)
+
+    def _launchRegular(self, A, options):
+        Job.Runner.startToil(A, options)
+
+    def _launchError(self, A, options):
+        try:
+            Job.Runner.startToil(A, options)
+        except FailedJobsException:
+            pass  # we expect a job to fail here
+        else:
+            self.fail("Toil run succeeded unexpectedly")
+
+def tempFileTestJob(job):
+    with open(job.fileStore.getLocalTempFile(), "w") as f:
+        f.write("test file retention")
+
+def tempFileTestErrorJob(job):
+    with open(job.fileStore.getLocalTempFile(), "w") as f:
+        f.write("test file retention")
+    raise RuntimeError()  # test failure