Repository: DataBiosphere/toil

Issue:
Add ability to hot-deploy entire virtualenvs
This is needed to support user distributions with dependencies. A user distribution is a distribution  with user modules. A distribution is a collection of modules, created with distutils/setuptools. The most prominent user distribution is toil-scripts. It may soon have a dependency (toil-lib, #922) and that will require this feature.

Add ability to hot-deploy entire virtualenvs
This is needed to support user distributions with dependencies. A user distribution is a distribution  with user modules. A distribution is a collection of modules, created with distutils/setuptools. The most prominent user distribution is toil-scripts. It may soon have a dependency (toil-lib, #922) and that will require this feature.



Discussion:

Resolving patch:
diff --git a/src/toil/__init__.py b/src/toil/__init__.py
--- a/src/toil/__init__.py
+++ b/src/toil/__init__.py
@@ -32,12 +32,16 @@ def toilPackageDirPath():
     return result
 
 
+def inVirtualEnv():
+    return hasattr(sys, 'real_prefix')
+
+
 def resolveEntryPoint(entryPoint):
     """
     Returns the path to the given entry point (see setup.py) that *should* work on a worker. The
     return value may be an absolute or a relative path.
     """
-    if hasattr(sys, 'real_prefix'):
+    if inVirtualEnv():
         path = os.path.join(os.path.dirname(sys.executable), entryPoint)
         # Inside a virtualenv we try to use absolute paths to the entrypoints. 
         if os.path.isfile(path):
diff --git a/src/toil/resource.py b/src/toil/resource.py
--- a/src/toil/resource.py
+++ b/src/toil/resource.py
@@ -11,24 +11,28 @@
 # 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
-from collections import namedtuple
-from contextlib import closing
+
+import errno
 import hashlib
 import importlib
-from io import BytesIO
 import json
 import logging
 import os
+import shutil
+import sys
+from collections import namedtuple
+from contextlib import closing
+from io import BytesIO
 from pydoc import locate
 from tempfile import mkdtemp
 from urllib2 import urlopen
-import errno
 from zipfile import ZipFile, PyZipFile
-import sys
-import shutil
+
 from bd2k.util.iterables import concat
-import tempfile
+
+from toil import inVirtualEnv
 
 log = logging.getLogger(__name__)
 
@@ -52,10 +56,10 @@ class Resource(namedtuple('Resource', ('name', 'pathHash', 'url', 'contentHash')
     ZIP archive of that directory.
     """
 
-    rootDirPathEnvName = 'JTRES_ROOT'
-
     resourceEnvNamePrefix = 'JTRES_'
 
+    rootDirPathEnvName = resourceEnvNamePrefix + 'ROOT'
+
     @classmethod
     def create(cls, jobStore, leaderPath):
         """
@@ -66,7 +70,10 @@ def create(cls, jobStore, leaderPath):
         :rtype: Resource
         """
         if os.path.isdir(leaderPath):
-            subcls = DirectoryResource
+            if inVirtualEnv() and leaderPath.startswith(sys.prefix):
+                subcls = VirtualEnvResource
+            else:
+                subcls = DirectoryResource
         elif os.path.isfile(leaderPath):
             subcls = FileResource
         elif os.path.exists(leaderPath):
@@ -98,7 +105,8 @@ def prepareSystem(cls):
         except KeyError:
             # Create directory holding local copies of requested resources ...
             resourceRootDirPath = mkdtemp()
-            # .. and register its location in an environment variable such that child processes can find it
+            # .. and register its location in an environment variable such that child processes
+            # can find it.
             os.environ[cls.rootDirPathEnvName] = resourceRootDirPath
         assert os.path.isdir(resourceRootDirPath)
 
@@ -108,10 +116,11 @@ def cleanSystem(cls):
         Removes all downloaded, localized resources
         """
         resourceRootDirPath = os.environ[cls.rootDirPathEnvName]
+        os.environ.pop(cls.rootDirPathEnvName)
         shutil.rmtree(resourceRootDirPath)
         for k, v in os.environ.items():
             if k.startswith(cls.resourceEnvNamePrefix):
-                os.unsetenv(k)
+                os.environ.pop(k)
 
     def register(self):
         """
@@ -166,8 +175,8 @@ def download(self, callback=None):
     @property
     def localPath(self):
         """
-        The path to resource on the worker. The file or directory at the given resource may not
-        exist. Invoking download() will ensure that it does.
+        The path to resource on the worker. The file or directory at the returned path may or may
+        not yet exist. Invoking download() will ensure that it does.
         """
         raise NotImplementedError
 
@@ -209,7 +218,7 @@ def _load(cls, path):
 
     def _save(self, dirPath):
         """
-        Save this resource to disk at the given parent path.
+        Save this resource to the directory at the given parent path.
 
         :type dirPath: str
         """
@@ -249,19 +258,20 @@ def localPath(self):
 class DirectoryResource(Resource):
     """
     A resource read from a directory on the leader. The URL will point to a ZIP archive of the
-    directory.
+    directory. Only Python script/modules will be included. The directory may be a package but it
+    does not need to be.
     """
 
     @classmethod
-    def _load(cls, leaderPath):
+    def _load(cls, path):
         """
-        :type leaderPath: str
+        :type path: str
         """
         bytesIO = BytesIO()
         # PyZipFile compiles .py files on the fly, filters out any non-Python files and
         # distinguishes between packages and simple directories.
         with PyZipFile(file=bytesIO, mode='w') as zipFile:
-            zipFile.writepy(leaderPath)
+            zipFile.writepy(path)
         bytesIO.seek(0)
         return bytesIO
 
@@ -277,6 +287,30 @@ def localPath(self):
         return self.localDirPath
 
 
+class VirtualEnvResource(DirectoryResource):
+    """
+    A resource read from a virtualenv on the leader. All modules and packages found in the
+    virtualenv's site-packages directory will be included. Any .pth or .egg-link files will be
+    ignored.
+    """
+
+    @classmethod
+    def _load(cls, path):
+        sitePackages = os.path.dirname(path)
+        assert os.path.basename(sitePackages) == 'site-packages'
+        bytesIO = BytesIO()
+        with PyZipFile(file=bytesIO, mode='w') as zipFile:
+            # This adds the .py files but omits subdirectories since site-packages is not a package
+            zipFile.writepy(sitePackages)
+            # Now add the missing packages
+            for name in os.listdir(sitePackages):
+                path = os.path.join(sitePackages, name)
+                if os.path.isdir(path) and os.path.isfile(os.path.join(path, '__init__.py')):
+                    zipFile.writepy(path)
+        bytesIO.seek(0)
+        return bytesIO
+
+
 class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name'))):
     """
     A path to a Python module decomposed into a namedtuple of three elements, namely
@@ -331,8 +365,9 @@ def forModule(cls, name):
         filePath[-1], extension = os.path.splitext(filePath[-1])
         assert extension in ('.py', '.pyc')
         if name == '__main__':
+            # User script/module was invoked as the main program
             if module.__package__:
-                # invoked via python -m foo.bar
+                # Invoked as a module via python -m foo.bar
                 name = [filePath.pop()]
                 for package in reversed(module.__package__.split('.')):
                     dirPathTail = filePath.pop()
@@ -341,17 +376,17 @@ def forModule(cls, name):
                 name = '.'.join(reversed(name))
                 dirPath = os.path.sep.join(filePath)
             else:
-                # invoked via python foo/bar.py
+                # Invoked as a script via python foo/bar.py
                 name = filePath.pop()
                 dirPath = os.path.sep.join(filePath)
                 cls._check_conflict(dirPath, name)
         else:
-            # imported as a module
+            # User module was imported. Determine the directory containing the top-level package
             for package in reversed(name.split('.')):
                 dirPathTail = filePath.pop()
                 assert dirPathTail == package
             dirPath = os.path.sep.join(filePath)
-
+        assert os.path.isdir(dirPath)
         return cls(dirPath=dirPath, name=name)
 
     @classmethod
@@ -402,7 +437,7 @@ def localize(self):
 
         :rtype: toil.resource.Resource
         """
-        if self._runningOnWorker():
+        if not self._runningOnWorker():
             log.warn('The localize() method should only be invoked on a worker.')
         resource = Resource.lookup(self._resourcePath)
         if resource is None:
@@ -441,8 +476,7 @@ def globalize(self):
             else:
                 raise
         else:
-            return self.__class__( dirPath=dirPath, name=self.name )
-
+            return self.__class__(dirPath=dirPath, name=self.name)
 
     @property
     def _resourcePath(self):


Test patch:
diff --git a/src/toil/test/src/resourceTest.py b/src/toil/test/src/resourceTest.py
--- a/src/toil/test/src/resourceTest.py
+++ b/src/toil/test/src/resourceTest.py
@@ -13,16 +13,18 @@
 # limitations under the License.
 
 from __future__ import absolute_import
+
 import importlib
 import os
-
 import sys
-from zipfile import ZipFile
-from bd2k.util.files import mkdir_p
 from io import BytesIO
+from subprocess import check_call
+from zipfile import ZipFile
 
+from bd2k.util.files import mkdir_p
 from mock import MagicMock, patch
 
+from toil import inVirtualEnv
 from toil.resource import ModuleDescriptor, Resource, ResourceException
 from toil.test import ToilTest
 
@@ -42,31 +44,59 @@ def testPackage(self):
                                                                  'foo/bar/__init__.py',
                                                                  'foo/bar/helper.py'))
 
+    def testVirtualEnv(self):
+        self._testExternal(moduleName='foo.userScript',
+                           virtualenv=True,
+                           pyFiles=('foo/__init__.py',
+                                    'foo/userScript.py',
+                                    'foo/bar/__init__.py',
+                                    'foo/bar/helper.py',
+                                    'de/pen/dency.py',
+                                    'de/__init__.py',
+                                    'de/pen/__init__.py'))
+
     def testStandAloneInPackage(self):
         self.assertRaises(ResourceException,
                           self._testExternal,
                           moduleName='userScript',
                           pyFiles=('__init__.py', 'userScript.py', 'helper.py'))
 
-    def _testExternal(self, moduleName, pyFiles):
+    def _testExternal(self, moduleName, pyFiles, virtualenv=False):
         dirPath = self._createTempDir()
-        pycFiles = set(pyFile + 'c' for pyFile in pyFiles)
-        for relPath in pyFiles:
-            path = os.path.join(dirPath, relPath)
-            mkdir_p(os.path.dirname(path))
-            with open(path, 'w') as f:
-                f.write('pass\n')
-        sys.path.append(dirPath)
+        if virtualenv:
+            self.assertTrue(inVirtualEnv())
+            check_call(['virtualenv', dirPath])
+            sitePackages = os.path.join(dirPath, 'lib', 'python2.7', 'site-packages')
+            # tuple assignment is necessary to make this line immediately precede the try:
+            oldPrefix, sys.prefix, dirPath = sys.prefix, dirPath, sitePackages
+        else:
+            oldPrefix = None
         try:
-            userScript = importlib.import_module(moduleName)
+            pycFiles = set(pyFile + 'c' for pyFile in pyFiles)
+            for relPath in pyFiles:
+                path = os.path.join(dirPath, relPath)
+                mkdir_p(os.path.dirname(path))
+                with open(path, 'w') as f:
+                    f.write('pass\n')
+            sys.path.append(dirPath)
             try:
-                self._test(userScript.__name__, expectedContents=pycFiles)
+                userScript = importlib.import_module(moduleName)
+                try:
+                    self._test(userScript.__name__,
+                               expectedContents=pycFiles,
+                               allowExtraContents=virtualenv)
+                finally:
+                    del userScript
+                    while moduleName:
+                        del sys.modules[moduleName]
+                        self.assertFalse(moduleName in sys.modules)
+                        moduleName = '.'.join(moduleName.split('.')[:-1])
+
             finally:
-                del userScript
-                del sys.modules[moduleName]
-            self.assertFalse(moduleName in sys.modules)
+                sys.path.remove(dirPath)
         finally:
-            sys.path.remove(dirPath)
+            if oldPrefix:
+                sys.prefix = oldPrefix
 
     def testBuiltIn(self):
         # Create a ModuleDescriptor for the module containing ModuleDescriptor, i.e. toil.resource
@@ -74,7 +104,8 @@ def testBuiltIn(self):
         self.assertEquals(module_name, 'toil.resource')
         self._test(module_name, shouldBelongToToil=True)
 
-    def _test(self, module_name, shouldBelongToToil=False, expectedContents=None):
+    def _test(self, module_name,
+              shouldBelongToToil=False, expectedContents=None, allowExtraContents=True):
         module = ModuleDescriptor.forModule(module_name)
         # Assert basic attributes and properties
         self.assertEqual(module.belongsToToil, shouldBelongToToil)
@@ -114,30 +145,37 @@ def _test(self, module_name, shouldBelongToToil=False, expectedContents=None):
         # Check contents if requested
         if expectedContents is not None:
             with ZipFile(BytesIO(zipFile)) as _zipFile:
-                self.assertEqual(set(_zipFile.namelist()), expectedContents)
+                actualContents = set(_zipFile.namelist())
+                if allowExtraContents:
+                    self.assertTrue(actualContents.issuperset(expectedContents))
+                else:
+                    self.assertEqual(actualContents, expectedContents)
 
         self.assertEquals(resource.url, url)
         # Now we're on the worker. Prepare the storage for localized resources
         Resource.prepareSystem()
-        # Register the resource for subsequent lookup.
-        resource.register()
-        # Lookup the resource and ensure that the result is equal to but not the same as the
-        # original resource. Lookup will also be used when we localize the module that was
-        # originally used to create the resource.
-        localResource = Resource.lookup(module._resourcePath)
-        self.assertEquals(resource, localResource)
-        self.assertIsNot(resource, localResource)
-        # Now show that we can localize the module using the registered resource. Set up a mock
-        # urlopen() that yields the zipped tree ...
-        mock_urlopen = MagicMock()
-        mock_urlopen.return_value.read.return_value = zipFile
-        with patch('toil.resource.urlopen', mock_urlopen):
-            # ... and use it to download and unpack the resource
-            localModule = module.localize()
-        # The name should be equal between original and localized resource ...
-        self.assertEquals(module.name, localModule.name)
-        # ... but the directory should be different.
-        self.assertNotEquals(module.dirPath, localModule.dirPath)
-        # Show that we can 'undo' localization. This is necessary when the user script's jobs are
-        # invoked on the worker where they generate more child jobs.
-        self.assertEquals(localModule.globalize(), module)
+        try:
+            # Register the resource for subsequent lookup.
+            resource.register()
+            # Lookup the resource and ensure that the result is equal to but not the same as the
+            # original resource. Lookup will also be used when we localize the module that was
+            # originally used to create the resource.
+            localResource = Resource.lookup(module._resourcePath)
+            self.assertEquals(resource, localResource)
+            self.assertIsNot(resource, localResource)
+            # Now show that we can localize the module using the registered resource. Set up a mock
+            # urlopen() that yields the zipped tree ...
+            mock_urlopen = MagicMock()
+            mock_urlopen.return_value.read.return_value = zipFile
+            with patch('toil.resource.urlopen', mock_urlopen):
+                # ... and use it to download and unpack the resource
+                localModule = module.localize()
+            # The name should be equal between original and localized resource ...
+            self.assertEquals(module.name, localModule.name)
+            # ... but the directory should be different.
+            self.assertNotEquals(module.dirPath, localModule.dirPath)
+            # Show that we can 'undo' localization. This is necessary when the user script's jobs
+            #  are invoked on the worker where they generate more child jobs.
+            self.assertEquals(localModule.globalize(), module)
+        finally:
+            Resource.cleanSystem()