Repository: DataBiosphere/toil

Issue:
Autoscaling does not work with restarted pipeline
After a pipeline is stopped & its nodes are terminated, rerunning the script with --restart will not spin up any more workers.



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
@@ -576,16 +576,7 @@ def start(self, rootJob):
             job = rootJob._serialiseFirstJob(self._jobStore)
             self._cacheJob(job)
 
-            if self.config.provisioner is None:
-                self._provisioner = None
-            elif self.config.provisioner == 'cgcloud':
-                logger.info('Using cgcloud provisioner.')
-                from toil.provisioners.cgcloud.provisioner import CGCloudProvisioner
-                self._provisioner = CGCloudProvisioner(self.config, self._batchSystem)
-            else:
-                # Command line parser shold have checked argument validity already
-                assert False, self.config.provisioner
-
+            self._setProvisioner()
             return self._runMainLoop(job)
         finally:
             self._shutdownBatchSystem()
@@ -607,11 +598,23 @@ def restart(self):
             self._setBatchSystemEnvVars()
             self._serialiseEnv()
             self._cacheAllJobs()
+            self._setProvisioner()
             rootJob = self._jobStore.clean(jobCache=self._jobCache)
             return self._runMainLoop(rootJob)
         finally:
             self._shutdownBatchSystem()
 
+    def _setProvisioner(self):
+        if self.config.provisioner is None:
+            self._provisioner = None
+        elif self.config.provisioner == 'cgcloud':
+            logger.info('Using cgcloud provisioner.')
+            from toil.provisioners.cgcloud.provisioner import CGCloudProvisioner
+            self._provisioner = CGCloudProvisioner(self.config, self._batchSystem)
+        else:
+            # Command line parser shold have checked argument validity already
+            assert False, self.config.provisioner
+
     @classmethod
     def getJobStore(cls, locator):
         """


Test patch:
diff --git a/run_tests.py b/run_tests.py
--- a/run_tests.py
+++ b/run_tests.py
@@ -28,7 +28,8 @@
         'GoogleJobStoreCacheTest',
         None],
     'integration-test': [
-        'CGCloudRNASeqTest'
+        'CGCloudRNASeqTest',
+        'CGCloudRestartTest'
     ]}
 
 
diff --git a/src/toil/test/__init__.py b/src/toil/test/__init__.py
--- a/src/toil/test/__init__.py
+++ b/src/toil/test/__init__.py
@@ -32,6 +32,7 @@
 import multiprocessing
 
 import time
+import signal
 from bd2k.util import less_strict_bool, memoize
 from bd2k.util.files import mkdir_p
 from bd2k.util.iterables import concat
@@ -446,6 +447,35 @@ def integrative(test_item):
 methodNamePartRegex = re.compile('^[a-zA-Z_0-9]+$')
 
 
+@contextmanager
+def timeLimit(seconds):
+    """
+    http://stackoverflow.com/a/601168
+    Use to limit the execution time of a function. Raises an exception if the execution of the
+    function takes more than the specified amount of time.
+
+    :param seconds: maximum allowable time, in seconds
+    >>> import time
+    >>> with timeLimit(5):
+    ...    time.sleep(4)
+    >>> import time
+    >>> with timeLimit(5):
+    ...    time.sleep(6)
+    Traceback (most recent call last):
+        ...
+    RuntimeError: Timed out
+    """
+    def signal_handler(signum, frame):
+        raise RuntimeError('Timed out')
+
+    signal.signal(signal.SIGALRM, signal_handler)
+    signal.alarm(seconds)
+    try:
+        yield
+    finally:
+        signal.alarm(0)
+
+
 # FIXME: move to bd2k-python-lib
 
 
diff --git a/src/toil/test/provisioners/cgcloud/provisionerTest.py b/src/toil/test/provisioners/cgcloud/provisionerTest.py
--- a/src/toil/test/provisioners/cgcloud/provisionerTest.py
+++ b/src/toil/test/provisioners/cgcloud/provisionerTest.py
@@ -16,15 +16,18 @@
 import os
 import pipes
 import subprocess
+
 from abc import abstractmethod, ABCMeta
+from inspect import getsource
+from textwrap import dedent
 from urlparse import urlparse
 from uuid import uuid4
 
 from bd2k.util.iterables import concat
 from cgcloud.lib.test import CgcloudTestCase
 
-from toil.test import integrative, ToilTest
-from toil.version import version as toil_version, cgcloudVersion
+from toil.test import integrative, ToilTest, timeLimit
+from toil.version import cgcloudVersion
 
 log = logging.getLogger(__name__)
 
@@ -199,9 +202,9 @@ def _runScript(self, toilOptions):
         raise NotImplementedError()
 
     @classmethod
-    def _cgcloud(cls, *args):
+    def _cgcloud(cls, *args, **kwargs):
         if not cls.dryRun:
-            cls._run('cgcloud', *args)
+            cls._run('cgcloud', *args, **kwargs)
 
     sshOptions = ['-o', 'UserKnownHostsFile=/dev/null', '-o', 'StrictHostKeyChecking=no']
 
@@ -215,7 +218,7 @@ def _ssh(cls, role, *args, **kwargs):
             del kwargs['admin']
 
         cls._cgcloud(
-            *filter(None, concat('ssh', '-a' if admin else None, role, cls.sshOptions, args)))
+            *filter(None, concat('ssh', '-a' if admin else None, role, cls.sshOptions, args)), **kwargs)
 
     @classmethod
     def _rsync(cls, role, *args):
@@ -265,3 +268,54 @@ def testAutoScaledCluster(self):
     @integrative
     def testAutoScaledSpotCluster(self):
         self._test(autoScaled=True, spotInstances=True)
+
+
+def restartScript():
+    from toil.job import Job
+    import argparse
+    import os
+
+    def f0(job):
+        if 'FAIL' in os.environ:
+            raise RuntimeError('failed on purpose')
+
+    if __name__ == '__main__':
+        parser = argparse.ArgumentParser()
+        Job.Runner.addToilOptions(parser)
+        options = parser.parse_args()
+        i = Job.Runner.startToil(Job.wrapJobFn(f0, cores=0.5, memory='50 M', disk='50 M'),
+                                 options)
+
+
+class CGCloudRestartTest(AbstractCGCloudProvisionerTest):
+    """
+    This test insures autoscaling works on a restarted Toil run
+    """
+    def setUp(self):
+        super(CGCloudRestartTest, self).setUp()
+        # CGCloud provisioner requires that the node has at least 1 ephemeral drive so this is the
+        # smallest instance we can use
+        self.instanceType = 'm3.medium'
+        self.leaderInstanceType = 'm3.medium'
+
+    def _getScript(self):
+        self.scriptName= "restartScript.py"
+        self._leader('tee', self.scriptName, input=dedent('\n'.join(getsource(restartScript).split('\n')[1:])),
+)
+
+    def _runScript(self, toilOptions):
+        # clean = onSuccess
+        disallowedOptions = ['--clean=always', '--retryCount=2']
+        newOptions = [option for option in toilOptions if option not in disallowedOptions]
+        try:
+            self._leader('python', self.scriptName, self.jobStore, '-e', 'FAIL=true', *newOptions)
+        except subprocess.CalledProcessError:
+            pass
+        else:
+            self.fail('Command succeeded when we expected failure')
+        with timeLimit(300):
+            self._leader('python', self.scriptName, self.jobStore, '--restart', *toilOptions)
+
+    @integrative
+    def testAutoScaledCluster(self):
+        self._test(autoScaled=True)