Repository: DataBiosphere/toil

Issue:
Toil context manager follow-up
remove job.FileStore from file managing heading in documentation

Toil context manager follow-up
remove job.FileStore from file managing heading in documentation



Discussion:
stress that all file staging is on the leader, and make contents show file staging as a subsection

In the example code, call exportFile() after run()

Also add a test that tests exactly that sequence: import, run, export

Have the job create an output file and return its ID such that you can export the output file rather than the input file. That would be a much more realistic scenario.

stress that all file staging is on the leader, and make contents show file staging as a subsection

In the example code, call exportFile() after run()

Also add a test that tests exactly that sequence: import, run, export

Have the job create an output file and return its ID such that you can export the output file rather than the input file. That would be a much more realistic scenario.


Resolving patch:
diff --git a/src/toil/common.py b/src/toil/common.py
--- a/src/toil/common.py
+++ b/src/toil/common.py
@@ -452,9 +452,10 @@ def __init__(self, options):
         super(Toil, self).__init__()
         self.options = options
         self.config = None
-        self.jobStore = None
-        self.batchSystem = None
-        self.jobCache = dict()
+        self._jobStore = None
+        self._batchSystem = None
+        self._jobCache = dict()
+        self._inContextManager = False
 
     def __enter__(self):
         """
@@ -462,15 +463,16 @@ def __enter__(self):
         consolidate the derived configuration with the one from the previous invocation of the
         workflow.
         """
+        self._inContextManager = True
         self.config = Config()
         self.config.setOptions(self.options)
-        self.jobStore = self.loadOrCreateJobStore(self.config.jobStore,
-                                                  config=None if self.config.restart else self.config)
+        self._jobStore = self.loadOrCreateJobStore(self.config.jobStore,
+                                                   config=None if self.config.restart else self.config)
         if self.config.restart:
             # Reload configuration from job store
-            self.config = self.jobStore.config
+            self.config = self._jobStore.config
             self.config.setOptions(self.options)
-            self.jobStore.writeConfigToStore()
+            self._jobStore.writeConfigToStore()
 
         return self
 
@@ -478,81 +480,78 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         """
         Clean up after a workflow invocaton. Depending on the configuration, delete the job store.
         """
+        self._inContextManager = False
         if (exc_type is not None and self.config.clean == "onError" or
             exc_type is None and self.config.clean == "onSuccess" or
-            self.config.clean == "always"):
-            self.jobStore.deleteJobStore()
+             self.config.clean == "always"):
+
+            self._jobStore.deleteJobStore()
         return False  # let exceptions through
 
-    def run(self, job=None):
+    def start(self, rootJob):
+        """
+        Invoke a Toil workflow with the given job as the root for an initial run. This method must
+        be called in the body of a ``with Toil(...) as toil:`` statement. This method should not be called
+        more than once for a workflow that has not finished.
+
+        :param toil.job.Job rootJob: The root job for the workflow
+        :return: The return value of the root job
+        :rtype: Any
+        """
+        self._assertContextManagerUsed()
+        if self.config.restart:
+            raise ToilConfigException('A workflow can only be started once. Use restart() to resume it.')
+
+        try:
+            self._batchSystem = self.createBatchSystem(self.config, jobStore=self._jobStore,
+                                                       userScript=rootJob.getUserScript())
+            self._setBatchSystemEnvVars()
+            self._serialiseEnv()
+            self._cacheAllJobs()
+
+            # Make a file to store the root job's return value in
+            rootJobReturnValueID = self._jobStore.getEmptyFileStoreID()
+
+            # Add the root job return value as a promise
+            if None not in rootJob._rvs:
+                rootJob._rvs[None] = []
+            rootJob._rvs[None].append(rootJobReturnValueID)
+
+            # Write the name of the promise file in a shared file
+            with self._jobStore.writeSharedFileStream("rootJobReturnValue") as fH:
+                fH.write(rootJobReturnValueID)
+
+            # Setup the first wrapper and cache it
+            job = rootJob._serialiseFirstJob(self._jobStore)
+            self._cacheJob(job)
+
+            return self._runMainLoop(job)
+
+        finally:
+            self._shutdownBatchSystem()
+
+    def restart(self):
         """
-        Invoke a Toil workflow with the given job as the root. The job argument may be omitted if
-        the invocaton is a restart in which case the previous invocation's root job will be used.
-        This method must be called in the body of a ``with Toil(...):`` statement.
+        Restarts a workflow that has been interrupted. This method should be called if and only if a workflow
+        has previously been started and has not finished.
 
-        :param toil.job.Job job: The root job for the workflow
         :return: The return value of the root job
         :rtype: Any
         """
-        self._assertContextManagerIsUsed()
-
-        userScript = None if job is None else job.getUserScript()
-        self.batchSystem = self.createBatchSystem(self.config,
-                                                  jobStore=self.jobStore,
-                                                  userScript=userScript)
-
-        realtimeLogLevel = self.options.logLevel if self.options.realTimeLogging else None
-        with RealtimeLogger(self.batchSystem, level=realtimeLogLevel):
-            try:
-                self._setBatchSystemEnvVars()
-                self._serialiseEnvironment()
-
-                logger.info('Caching all jobs in job store')
-                self._cacheJobs()
-                logger.info('{} jobs downloaded.'.format(len(self.jobCache)))
-
-                if job is not None:
-                    if self.config.restart:
-                        raise ToilConfigException('Cannot override root job on restart')
-
-                    # Make a file to store the root job's return value in
-                    jobStoreFileID = self.jobStore.getEmptyFileStoreID()
-
-                    # Add the root job return value as a promise
-                    if None not in job._rvs:
-                        job._rvs[None] = []
-                    job._rvs[None].append(jobStoreFileID)
-
-                    # Write the name of the promise file in a shared file
-                    with self.jobStore.writeSharedFileStream("rootJobReturnValue") as fH:
-                        fH.write(jobStoreFileID)
-
-                    # Setup the first wrapper and cache it
-                    rootJob = job._serialiseFirstJob(self.jobStore)
-                    self._cacheJob(rootJob)
-
-                else:  # restarting
-                    if not self.config.restart:
-                        raise ToilConfigException('Need root job on first invocation')
-
-                    # This cleans up any half written jobs after a restart
-                    rootJob = self.jobStore.clean(jobCache=self.jobCache)
-
-                # FIXME: common.py shouldn't import from leader.py
-                from toil.leader import mainLoop
-                return mainLoop(config=self.config,
-                                batchSystem=self.batchSystem,
-                                provisioner=None,
-                                jobStore=self.jobStore,
-                                rootJobWrapper=rootJob,
-                                jobCache=self.jobCache)
-
-            finally:
-                startTime = time.time()
-                logger.debug('Shutting down batch system')
-                self.batchSystem.shutdown()
-                logger.debug('Finished shutting down the batch system in %s seconds.' % (time.time() - startTime))
+        self._assertContextManagerUsed()
+        if not self.config.restart:
+            raise ToilConfigException('Cannot call restart on initial run of workflow')
 
+        try:
+            self._batchSystem = self.createBatchSystem(self.config, jobStore=self._jobStore)
+            self._setBatchSystemEnvVars()
+            self._serialiseEnv()
+            self._cacheAllJobs()
+
+            rootJob = self._jobStore.clean(jobCache=self._jobCache)
+            return self._runMainLoop(rootJob)
+        finally:
+            self._shutdownBatchSystem()
 
     @staticmethod
     def loadOrCreateJobStore(jobStoreString, config=None):
@@ -651,43 +650,45 @@ def createBatchSystem(config, jobStore=None, userScript=None):
         return batchSystemClass(**kwargs)
 
     def importFile(self, srcUrl):
-        self._assertContextManagerIsUsed()
-        return self.jobStore.importFile(srcUrl)
+        self._assertContextManagerUsed()
+        return self._jobStore.importFile(srcUrl)
 
     def exportFile(self, jobStoreFileID, dstUrl):
-        self._assertContextManagerIsUsed()
-        self.jobStore.exportFile(jobStoreFileID, dstUrl)
+        self._assertContextManagerUsed()
+        self._jobStore.exportFile(jobStoreFileID, dstUrl)
 
     def _setBatchSystemEnvVars(self):
         """
         Sets the environment variables required by the job store and those passed on command line.
         """
-        for envDict in (self.jobStore.getEnv(), self.config.environment):
+        for envDict in (self._jobStore.getEnv(), self.config.environment):
             for k, v in envDict.iteritems():
-                self.batchSystem.setEnv(k, v)
+                self._batchSystem.setEnv(k, v)
 
-    def _serialiseEnvironment(self):
+    def _serialiseEnv(self):
         """
         Puts the environment in a globally accessible pickle file.
         """
         # Dump out the environment of this process in the environment pickle file.
-        with self.jobStore.writeSharedFileStream("environment.pickle") as fileHandle:
+        with self._jobStore.writeSharedFileStream("environment.pickle") as fileHandle:
             cPickle.dump(os.environ, fileHandle, cPickle.HIGHEST_PROTOCOL)
         logger.info("Written the environment for the jobs to the environment file")
 
-    def _cacheJobs(self):
+    def _cacheAllJobs(self):
         """
         Downloads all jobs in the current job store into self.jobCache.
         """
-        self.jobCache = {jobWrapper.jobStoreID: jobWrapper for jobWrapper in self.jobStore.jobs()}
+        logger.info('Caching all jobs in job store')
+        self._jobCache = {jobWrapper.jobStoreID: jobWrapper for jobWrapper in self._jobStore.jobs()}
+        logger.info('{} jobs downloaded.'.format(len(self._jobCache)))
 
     def _cacheJob(self, job):
         """
-        Adds given job to curent jobCache.
+        Adds given job to current job cache.
 
         :param toil.jobWrapper.JobWrapper job: job to be added to current job cache
         """
-        self.jobCache[job.jobStoreID] = job
+        self._jobCache[job.jobStoreID] = job
 
     @staticmethod
     def getWorkflowDir(workflowID, configWorkDir=None):
@@ -717,11 +718,36 @@ def getWorkflowDir(workflowID, configWorkDir=None):
             logger.info('Created the workflow directory at %s' % workflowDir)
         return workflowDir
 
-    def _assertContextManagerIsUsed(self):
-        # Assert that we are inside the context manager
-        if self.jobStore is None:
-            # __enter__() sets self.jobStore
-            raise ToilContextManagerMisuseException("Toil class cannot be instantiated outside of a context manager")
+    def _runMainLoop(self, rootJob):
+        """
+        Runs the main loop with the given job.
+        :param toil.job.Job rootJob: The root job for the workflow.
+        :rtype: Any
+        """
+        with RealtimeLogger(self._batchSystem, level=self.options.logLevel if self.options.realTimeLogging else None):
+            # FIXME: common should not import from leader
+            from toil.leader import mainLoop
+            return mainLoop(config=self.config,
+                            batchSystem=self._batchSystem,
+                            provisioner=None,
+                            jobStore=self._jobStore,
+                            rootJobWrapper=rootJob,
+                            jobCache=self._jobCache)
+
+    def _shutdownBatchSystem(self):
+        """
+        Shuts down current batch system if it has been created.
+        """
+        assert self._batchSystem is not None
+
+        startTime = time.time()
+        logger.debug('Shutting down batch system')
+        self._batchSystem.shutdown()
+        logger.debug('Finished shutting down the batch system in %s seconds.' % (time.time() - startTime))
+
+    def _assertContextManagerUsed(self):
+        if not self._inContextManager:
+            raise ToilContextManagerMisuseException("This method cannot be called outside of Toil context manager.")
 
 
 class ToilConfigException(Exception):
@@ -743,6 +769,7 @@ def __init__(self, message):
 
 # Nested functions can't have doctests so we have to make this global
 
+
 def parseSetEnv(l):
     """
     Parses a list of strings of the form "NAME=VALUE" or just "NAME" into a dictionary. Strings
diff --git a/src/toil/cwl/cwltoil.py b/src/toil/cwl/cwltoil.py
--- a/src/toil/cwl/cwltoil.py
+++ b/src/toil/cwl/cwltoil.py
@@ -582,7 +582,7 @@ def importDefault(tool):
         adjustFiles(builder.job, functools.partial(writeFile, toil.importFile, {}))
         wf1.cwljob = builder.job
 
-        outobj = toil.run(wf1)
+        outobj = toil.start(wf1)
         outobj = resolve_indirect(outobj)
 
         adjustFilesWithSecondary(outobj, functools.partial(getFile, toil, outdir, index={}, export=True, rename_collision=True))
diff --git a/src/toil/job.py b/src/toil/job.py
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -474,10 +474,10 @@ def startToil(job, options):
             """
             setLoggingFromOptions(options)
             with Toil(options) as toil:
-                if options.restart:
-                    job = None
-                return toil.run(job)
-
+                if not options.restart:
+                    return toil.start(job)
+                else:
+                    return toil.restart()
 
     class FileStore( object ):
         """


Test patch:
diff --git a/src/toil/test/sort/sortTest.py b/src/toil/test/sort/sortTest.py
--- a/src/toil/test/sort/sortTest.py
+++ b/src/toil/test/sort/sortTest.py
@@ -22,7 +22,7 @@
 from toil import resolveEntryPoint
 
 from toil.batchSystems.parasolTestSupport import ParasolTestSupport
-from toil.common import Toil
+from toil.common import Toil, ToilConfigException
 from toil.job import Job, JobException
 from toil.lib.bioio import getLogLevelString
 from toil.batchSystems.mesos.test import MesosTestSupport
@@ -96,7 +96,7 @@ def _toilSort(self, jobStore, batchSystem,
                 options.restart = True
                 try:
                     with Toil(options) as toil:
-                        toil.run()
+                        toil.restart()
                     self.fail()
                 except JobStoreCreationException:
                     pass
@@ -106,7 +106,7 @@ def _toilSort(self, jobStore, batchSystem,
                 # Now actually run the workflow
                 try:
                     with Toil(options) as toil:
-                        toil.run(firstJob)
+                        toil.start(firstJob)
                     i = 0
                 except FailedJobsException as e:
                     i = e.numberOfFailedJobs
@@ -114,7 +114,7 @@ def _toilSort(self, jobStore, batchSystem,
                 # Check we get an exception if we try to run without restart on an existing store
                 try:
                     with Toil(options) as toil:
-                        toil.run(firstJob)
+                        toil.start(firstJob)
                     self.fail()
                 except JobStoreCreationException:
                     pass
@@ -127,7 +127,7 @@ def _toilSort(self, jobStore, batchSystem,
                     options.useExistingOptions = random.random() > 0.5
                     try:
                         with Toil(options) as toil:
-                            toil.run()
+                            toil.restart()
                         i = 0
                     except FailedJobsException as e:
                         i = e.numberOfFailedJobs
@@ -139,7 +139,7 @@ def _toilSort(self, jobStore, batchSystem,
                 # indicating that there are no jobs remaining in the workflow.
                 try:
                     with Toil(options) as toil:
-                        toil.run()
+                        toil.restart()
                 except JobException:
                     pass
                 else:
diff --git a/src/toil/test/src/importExportFileTest.py b/src/toil/test/src/importExportFileTest.py
--- a/src/toil/test/src/importExportFileTest.py
+++ b/src/toil/test/src/importExportFileTest.py
@@ -16,62 +16,75 @@
 
 import uuid
 
-import shutil
-
-import os
 from toil.common import Toil
 from toil.job import Job
-from toil.test import ToilTest
-import tempfile
+from toil.leader import FailedJobsException
+from toil.test import ToilTest, make_tests
 
 
 class ImportExportFileTest(ToilTest):
     def setUp(self):
-        self.srcFile = "%s/%s" % (tempfile.mkdtemp(), uuid.uuid4())
-        with open(self.srcFile, 'w') as f:
-            f.write(os.urandom(2**10))
-
-    def tearDown(self):
-        shutil.rmtree(self.srcFile.rsplit('/', 1)[0])
+        super(ImportExportFileTest, self).setUp()
+        self._tempDir = self._createTempDir()
+        self.dstFile = '%s/%s' % (self._tempDir, 'out')
 
-    def testImportExportFile(self):
-        options = Job.Runner.getDefaultOptions(self._getTestJobStorePath())
-        options.logLevel = "INFO"
-        job = HelloWorld()
+    def _importExportFile(self, options, fail):
         with Toil(options) as toil:
-            jobStoreFileID = toil.importFile('file://%s' % self.srcFile)
-            assert toil.jobStore.fileExists(jobStoreFileID)
+            if not options.restart:
 
-            destPath = '%s/%s' % (self.srcFile.rsplit('/', 1)[0], uuid.uuid4())
-            toil.exportFile(jobStoreFileID, 'file://%s' % destPath)
-            assert os.path.exists(destPath)
+                srcFile = '%s/%s%s' % (self._tempDir, 'in', uuid.uuid4())
+                with open(srcFile, 'w') as f:
+                    f.write('Hello')
+                inputFileID = toil.importFile('file://' + srcFile)
 
-            toil.run(job)
+                # Write a boolean that determines whether the job fails.
+                with toil._jobStore.writeFileStream() as (f, failFileID):
+                    self.failFileID = failFileID
+                    f.write(str(fail))
 
+                outputFileID = toil.start(HelloWorld(inputFileID, self.failFileID))
+            else:
+                # Set up job for failure
+                with toil._jobStore.updateFileStream(self.failFileID) as f:
+                    f.write('False')
 
-class HelloWorld(Job):
-    def __init__(self):
-        Job.__init__(self,  memory=100000, cores=2, disk="3G")
+                outputFileID = toil.restart()
 
-    def run(self, fileStore):
-        fileID = self.addChildJobFn(childFn, cores=1, memory="1M", disk="1M").rv()
-        self.addFollowOn(FollowOn(fileID))
+            toil.exportFile(outputFileID, 'file://' + self.dstFile)
+            with open(self.dstFile, 'r') as f:
+                assert f.read() == "HelloWorld!"
 
+    def _importExport(self, restart):
+        options = Job.Runner.getDefaultOptions(self._getTestJobStorePath())
+        options.logLevel = "INFO"
+
+        if restart:
+            try:
+                self._importExportFile(options, fail=True)
+            except FailedJobsException:
+                options.restart = True
 
-def childFn(job):
-    with job.fileStore.writeGlobalFileStream() as (fH, fileID):
-        fH.write("Hello, World!")
-        return fileID
+        self._importExportFile(options, fail=False)
 
+    def testImportExportRestartTrue(self):
+        self._importExport(restart=True)
 
-class FollowOn(Job):
-    def __init__(self,fileId):
-        Job.__init__(self)
-        self.fileId=fileId
+    def testImportExportRestartFalse(self):
+        self._importExport(restart=False)
+
+
+class HelloWorld(Job):
+    def __init__(self, inputFileID, failFileID):
+        Job.__init__(self,  memory=100000, cores=1, disk="1M")
+        self.inputFileID = inputFileID
+        self.failFileID = failFileID
 
     def run(self, fileStore):
-        tempDir = fileStore.getLocalTempDir()
-        tempFilePath = "/".join([tempDir,"LocalCopy"])
-        with fileStore.readGlobalFileStream(self.fileId) as globalFile:
-            with open(tempFilePath, "w") as localFile:
-                localFile.write(globalFile.read())
+        with fileStore.readGlobalFileStream(self.failFileID) as failValue:
+            if failValue.read() == 'True':
+                raise RuntimeError('planned exception')
+            else:
+                with fileStore.readGlobalFileStream(self.inputFileID) as fi:
+                    with fileStore.writeGlobalFileStream() as (fo, outputFileID):
+                        fo.write(fi.read() + 'World!')
+                        return outputFileID
diff --git a/src/toil/test/src/jobWrapperTest.py b/src/toil/test/src/jobWrapperTest.py
--- a/src/toil/test/src/jobWrapperTest.py
+++ b/src/toil/test/src/jobWrapperTest.py
@@ -33,7 +33,7 @@ def setUp(self):
 
     def tearDown(self):
         self.toil.__exit__(None, None, None)
-        self.toil.jobStore.deleteJobStore()
+        self.toil._jobStore.deleteJobStore()
         self.assertFalse(os.path.exists(self.jobStorePath))
         super(JobWrapperTest, self).tearDown()
     
diff --git a/src/toil/test/src/toilContextManagerTest.py b/src/toil/test/src/toilContextManagerTest.py
--- a/src/toil/test/src/toilContextManagerTest.py
+++ b/src/toil/test/src/toilContextManagerTest.py
@@ -23,13 +23,13 @@ def testContextManger(self):
         options = Job.Runner.getDefaultOptions(self._getTestJobStorePath())
         options.logLevel = 'INFO'
         with Toil(options) as toil:
-            toil.run(HelloWorld())
+            toil.start(HelloWorld())
 
     def testNoContextManger(self):
         options = Job.Runner.getDefaultOptions(self._getTestJobStorePath())
         options.logLevel = 'INFO'
         toil = Toil(options)
-        self.assertRaises(ToilContextManagerMisuseException, toil.run, HelloWorld())
+        self.assertRaises(ToilContextManagerMisuseException, toil.start, HelloWorld())
 
 
 class HelloWorld(Job):