Repository: databiosphere/toil
Issue: Hot-deployment fails with job functions spread over two top-level packages

The following structure in a venv on the leader /path/to/virtualenv/lib/python/site-packages/toil_scripts
/path/to/virtualenv/lib/python/site-packages/toil_lib

causes hot-deployment to fail localization in a job function if parent job function is from toil_scripts.foo.bla while the other job function is from toil_lib.bar.fasel.


Changed files:
src/toil/test/src/hotDeploymentTest.py, src/toil/resource.py, src/toil/test/__init__.py, src/toil/job.py

diff --git a/src/toil/job.py b/src/toil/job.py
index 651b24db..a0586214 100644
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -690,7 +690,7 @@ class Job(object):
         """
         commandTokens = command.split()
         assert "_toil" == commandTokens[0]
-        userModule = ModuleDescriptor(*(commandTokens[2:]))
+        userModule = ModuleDescriptor.fromCommand(commandTokens[2:])
         logger.debug('Loading user module %s.', userModule)
         userModule = cls._loadUserModule(userModule)
         pickleFile = commandTokens[1]
@@ -909,7 +909,7 @@ class Job(object):
         # and FunctionWrappingJob overrides getUserScript() to give us just that. Only then can
         # filter_main() in _unpickle( ) do its job of resolving any user-defined type or function.
         userScript = self.getUserScript().globalize()
-        jobsToJobWrappers[self].command = ' '.join( ('_toil', fileStoreID) + userScript)
+        jobsToJobWrappers[self].command = ' '.join(('_toil', fileStoreID) + userScript.toCommand())
         #Update the status of the jobWrapper on disk
         jobStore.update(jobsToJobWrappers[self])
 
diff --git a/src/toil/resource.py b/src/toil/resource.py
index 6ca0eb6b..bbeb30a0 100644
--- a/src/toil/resource.py
+++ b/src/toil/resource.py
@@ -30,6 +30,7 @@ from tempfile import mkdtemp
 from urllib2 import urlopen
 from zipfile import ZipFile, PyZipFile
 
+from bd2k.util import strict_bool
 from bd2k.util.iterables import concat
 from bd2k.util.exceptions import require
 
@@ -70,30 +71,18 @@ class Resource(namedtuple('Resource', ('name', 'pathHash', 'url', 'contentHash')
 
         :rtype: Resource
         """
-        if os.path.isdir(leaderPath):
-            if inVirtualEnv() and leaderPath.startswith(sys.prefix):
-                subcls = VirtualEnvResource
-            else:
-                subcls = DirectoryResource
-        elif os.path.isfile(leaderPath):
-            subcls = FileResource
-        elif os.path.exists(leaderPath):
-            raise AssertionError("Neither a file or a directory: '%s'" % leaderPath)
-        else:
-            raise AssertionError("No such file or directory: '%s'" % leaderPath)
-
-        pathHash = subcls._pathHash(leaderPath)
+        pathHash = cls._pathHash(leaderPath)
         contentHash = hashlib.md5()
         # noinspection PyProtectedMember
-        with subcls._load(leaderPath) as src:
+        with cls._load(leaderPath) as src:
             with jobStore.writeSharedFileStream(pathHash, isProtected=False) as dst:
                 userScript = src.read()
                 contentHash.update(userScript)
                 dst.write(userScript)
-        return subcls(name=os.path.basename(leaderPath),
-                      pathHash=pathHash,
-                      url=(jobStore.getSharedPublicUrl(pathHash)),
-                      contentHash=contentHash.hexdigest())
+        return cls(name=os.path.basename(leaderPath),
+                   pathHash=pathHash,
+                   url=(jobStore.getSharedPublicUrl(pathHash)),
+                   contentHash=contentHash.hexdigest())
 
     @classmethod
     def prepareSystem(cls):
@@ -297,7 +286,7 @@ class VirtualEnvResource(DirectoryResource):
 
     @classmethod
     def _load(cls, path):
-        sitePackages = os.path.dirname(path)
+        sitePackages = path
         assert os.path.basename(sitePackages) == 'site-packages'
         bytesIO = BytesIO()
         with PyZipFile(file=bytesIO, mode='w') as zipFile:
@@ -312,7 +301,7 @@ class VirtualEnvResource(DirectoryResource):
         return bytesIO
 
 
-class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
+class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name', 'fromVirtualEnv'))):
     """
     A path to a Python module decomposed into a namedtuple of three elements, namely
 
@@ -324,7 +313,7 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
 
     >>> import toil.resource
     >>> ModuleDescriptor.forModule('toil.resource') # doctest: +ELLIPSIS
-    ModuleDescriptor(dirPath='/.../src', name='toil.resource')
+    ModuleDescriptor(dirPath='/.../src', name='toil.resource', fromVirtualEnv=False)
 
     >>> import subprocess, tempfile, os
     >>> dirPath = tempfile.mkdtemp()
@@ -333,7 +322,7 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
     ...     f.write('from toil.resource import ModuleDescriptor\\\\n'
     ...             'print ModuleDescriptor.forModule(__name__)')
     >>> subprocess.check_output([ sys.executable, path ]) # doctest: +ELLIPSIS
-    "ModuleDescriptor(dirPath='...', name='foo')\\\\n"
+    "ModuleDescriptor(dirPath='...', name='foo', fromVirtualEnv=False)\\\\n"
 
     Now test a collision. As funny as it sounds, the robotparser module is included in the Python
     standard library.
@@ -389,7 +378,8 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
                 assert dirPathTail == package
             dirPath = os.path.sep.join(filePath)
         assert os.path.isdir(dirPath)
-        return cls(dirPath=dirPath, name=name)
+        fromVirtualEnv = inVirtualEnv() and dirPath.startswith(sys.prefix)
+        return cls(dirPath=dirPath, name=name, fromVirtualEnv=fromVirtualEnv)
 
     @classmethod
     def _check_conflict(cls, dirPath, name):
@@ -429,7 +419,23 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
         :type jobStore: toil.jobStores.abstractJobStore.AbstractJobStore
         :rtype: toil.resource.Resource
         """
-        return Resource.create(jobStore, self._resourcePath)
+        return self._getResourceClass().create(jobStore, self._resourcePath)
+
+    def _getResourceClass(self):
+        """
+        Return the concrete subclass of Resource that's appropriate for hot-deploying this module.
+        """
+        if self.fromVirtualEnv:
+            subcls = VirtualEnvResource
+        elif os.path.isdir(self._resourcePath):
+            subcls = DirectoryResource
+        elif os.path.isfile(self._resourcePath):
+            subcls = FileResource
+        elif os.path.exists(self._resourcePath):
+            raise AssertionError("Neither a file or a directory: '%s'" % self._resourcePath)
+        else:
+            raise AssertionError("No such file or directory: '%s'" % self._resourcePath)
+        return subcls
 
     def localize(self):
         """
@@ -448,11 +454,14 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
         else:
             def stash(tmpDirPath):
                 # Save the original dirPath such that we can restore it in globalize()
-                with open(os.path.join(tmpDirPath, '.dirPath'), 'w') as f:
+                with open(os.path.join(tmpDirPath, '.stash'), 'w') as f:
+                    f.write('1' if self.fromVirtualEnv else '0')
                     f.write(self.dirPath)
 
             resource.download(callback=stash)
-            return self.__class__(dirPath=resource.localDirPath, name=self.name)
+            return self.__class__(dirPath=resource.localDirPath,
+                                  name=self.name,
+                                  fromVirtualEnv=self.fromVirtualEnv)
 
     def _runningOnWorker(self):
         try:
@@ -471,7 +480,8 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
         Reverse the effect of localize().
         """
         try:
-            with open(os.path.join(self.dirPath, '.dirPath')) as f:
+            with open(os.path.join(self.dirPath, '.stash')) as f:
+                fromVirtualEnv = [False, True][int(f.read(1))]
                 dirPath = f.read()
         except IOError as e:
             if e.errno == errno.ENOENT:
@@ -481,7 +491,9 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
             else:
                 raise
         else:
-            return self.__class__(dirPath=dirPath, name=self.name)
+            return self.__class__(dirPath=dirPath,
+                                  name=self.name,
+                                  fromVirtualEnv=fromVirtualEnv)
 
     @property
     def _resourcePath(self):
@@ -489,7 +501,9 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
         The path to the directory that should be used when shipping this module and its siblings
         around as a resource.
         """
-        if '.' in self.name:
+        if self.fromVirtualEnv:
+            return self.dirPath
+        elif '.' in self.name:
             return os.path.join(self.dirPath, self._rootPackage())
         else:
             initName = self._initModuleName(self.dirPath)
@@ -519,6 +533,14 @@ class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
         else:
             return head
 
+    def toCommand(self):
+        return tuple(map(str,self))
+
+    @classmethod
+    def fromCommand(cls, command):
+        assert len(command) == 3
+        return cls( dirPath=command[0], name=command[1], fromVirtualEnv=strict_bool(command[2]))
+
 
 class ResourceException(Exception):
     pass
diff --git a/src/toil/test/__init__.py b/src/toil/test/__init__.py
index c0a05d68..64bb8581 100644
--- a/src/toil/test/__init__.py
+++ b/src/toil/test/__init__.py
@@ -24,6 +24,8 @@ import re
 import uuid
 from abc import ABCMeta, abstractmethod
 from contextlib import contextmanager
+from inspect import getsource
+from textwrap import dedent
 from unittest.util import strclass
 from urllib2 import urlopen
 
@@ -713,6 +715,30 @@ class ApplianceTestSupport(ToilTest):
         def writeToAppliance(self, path, contents):
             self.runOnAppliance('tee', path, input=contents)
 
+        def deployScript(self, path, packagePath, script):
+            """
+            Deploy a Python module on the appliance.
+
+            :param path: the path (absolute or relative to the WORDIR of the appliance container)
+                   to the root of the package hierarchy where the given module should be placed.
+                   The given directory should be on the Python path.
+
+            :param packagePath: the desired fully qualified module name (dotted form) of the module
+
+            :param str|callable script: the contents of the Python module. If a callable is given,
+                   its source code will be extracted. This is a convenience that lets you embed
+                   user scripts into test code as nested function.
+            """
+            if callable(script):
+                script = dedent('\\n'.join(getsource(script).split('\\n')[1:]))
+            packagePath = packagePath.split('.')
+            packages, module = packagePath[:-1], packagePath[-1]
+            for package in packages:
+                path += '/' + package
+                self.runOnAppliance('mkdir', path)
+                self.writeToAppliance(path + '/__init__.py', '')
+            self.writeToAppliance(path + '/' + module + '.py', script)
+
     class LeaderThread(Appliance):
         def _getRole(self):
             return 'leader'
diff --git a/src/toil/test/src/hotDeploymentTest.py b/src/toil/test/src/hotDeploymentTest.py
index 33d80e82..41227876 100644
--- a/src/toil/test/src/hotDeploymentTest.py
+++ b/src/toil/test/src/hotDeploymentTest.py
@@ -1,7 +1,6 @@
 import logging
-from inspect import getsource
+from contextlib import contextmanager
 from subprocess import CalledProcessError
-from textwrap import dedent
 
 from bd2k.util.iterables import concat
 
@@ -16,15 +15,32 @@ class HotDeploymentTest(ApplianceTestSupport):
         logging.basicConfig(level=logging.INFO)
         super(HotDeploymentTest, self).setUp()
 
-    def test(self):
+    @contextmanager
+    def _venvApplianceCluster(self):
+        """
+        Creates an appliance cluster with a virtualenv at './venv' on the leader and a temporary
+        directory on the host mounted at /data in the leader and worker containers.
+        """
         dataDirPath = self._createTempDir(purpose='data')
         with self._applianceCluster(mounts={dataDirPath: '/data'}) as (leader, worker):
-            leader.runOnAppliance('virtualenv',
-                                  '--system-site-packages',
-                                  '--never-download',
-                                  'venv')
-            leader.runOnAppliance('venv/bin/pip', 'list')
+            try:
+                leader.runOnAppliance('virtualenv',
+                                      '--system-site-packages',
+                                      '--never-download',
+                                      'venv')
+                leader.runOnAppliance('venv/bin/pip', 'list') # For diagnostic purposes
+                yield leader, worker
+            finally:
+                # Without this step, we would leak files owned by root on the host's file system
+                leader.runOnAppliance('rm', '-rf', '/data/*')
 
+    sitePackages = 'venv/lib/python2.7/site-packages'
+
+    def testRestart(self):
+        """
+        Test whether hot-deployment works on restart.
+        """
+        with self._venvApplianceCluster() as (leader, worker):
             def userScript():
                 from toil.job import Job
                 from toil.common import Toil
@@ -42,11 +58,10 @@ class HotDeploymentTest(ApplianceTestSupport):
                         else:
                             toil.start(Job.wrapJobFn(job))
 
-            package = 'venv/lib/python2.7/site-packages/foo'
-            leader.runOnAppliance('mkdir', '-p', package)
-            leader.writeToAppliance(package + '/__init__.py', '')
-            script = dedent('\\n'.join(getsource(userScript).split('\\n')[1:]))
-            leader.writeToAppliance(package + '/bar.py', script)
+            leader.deployScript(path=self.sitePackages,
+                                packagePath='foo.bar',
+                                script=userScript)
+
             pythonArgs = ['venv/bin/python', '-m', 'foo.bar']
             toilArgs = ['--logDebug',
                         '--batchSystem=mesos',
@@ -57,3 +72,64 @@ class HotDeploymentTest(ApplianceTestSupport):
             self.assertRaises(CalledProcessError, leader.runOnAppliance, *command)
             command = concat(pythonArgs, '--restart', toilArgs)
             leader.runOnAppliance(*command)
+
+    def testSplitRootPackages(self):
+        """
+        Test whether hot-deployment works with a virtualenv in which jobs are defined in
+        completely separate branches of the package hierarchy. Initially, hot deployment did
+        deploy the entire virtualenv but jobs could only be defined in one branch of the package
+        hierarchy. We define a branch as the maximum set of fully qualified package paths that
+        share the same first component. IOW, a.b and a.c are in the same branch, while a.b and
+        d.c are not.
+        """
+        with self._venvApplianceCluster() as (leader, worker):
+
+            # Deploy the library module with job definitions
+            def libraryModule():
+                # noinspection PyUnusedLocal
+                def libraryJob(job):
+                    open('/data/foo.txt', 'w').close()
+
+            leader.deployScript(path=self.sitePackages,
+                                packagePath='toil_lib.foo',
+                                script=libraryModule)
+
+            # Deploy the user script
+            def userScript():
+                from toil.job import Job
+                from toil.common import Toil
+                # noinspection PyUnresolvedReferences
+                from toil_lib.foo import libraryJob
+
+                # noinspection PyUnusedLocal
+                def job(job, disk='10M', cores=1, memory='10M'):
+                    # Double the requirements to prevent chaining as chaining might hide problems
+                    # in hot deployment code.
+                    job.addChildJobFn(libraryJob, disk='20M', cores=cores, memory=memory)
+
+                if __name__ == '__main__':
+                    options = Job.Runner.getDefaultArgumentParser().parse_args()
+                    with Toil(options) as toil:
+                        if toil.config.restart:
+                            toil.restart()
+                        else:
+                            toil.start(Job.wrapJobFn(job))
+
+            leader.deployScript(path=self.sitePackages,
+                                packagePath='toil_script.bar',
+                                script=userScript)
+
+            # Assert that output file isn't there
+            worker.runOnAppliance('test', '!', '-f', '/data/foo.txt')
+            # Just being paranoid
+            self.assertRaises(CalledProcessError,
+                              worker.runOnAppliance, 'test', '-f', '/data/foo.txt')
+            leader.runOnAppliance('venv/bin/python',
+                                  '-m', 'toil_script.bar',
+                                  '--logDebug',
+                                  '--batchSystem=mesos',
+                                  '--mesosMaster=localhost:5050',
+                                  '--defaultMemory=10M',
+                                  '/data/jobstore')
+            # Assert that out output file is there
+            worker.runOnAppliance('test', '-f', '/data/foo.txt')