Repository: DataBiosphere/toil

Issue:
Clean up resource directory on executor shutdown



Discussion:

Resolving patch:
diff --git a/src/toil/batchSystems/mesos/executor.py b/src/toil/batchSystems/mesos/executor.py
--- a/src/toil/batchSystems/mesos/executor.py
+++ b/src/toil/batchSystems/mesos/executor.py
@@ -60,23 +60,24 @@ def reregistered(self, driver, slaveInfo):
         """
         Invoked when the executor re-registers with a restarted slave.
         """
-        log.info("Re-Registered")
+        log.info("Re-registered")
 
     def disconnected(self, driver):
         """
         Invoked when the executor becomes "disconnected" from the slave (e.g., the slave is being restarted due to an upgrade).
         """
-        log.critical("disconnected from slave")
+        log.critical("Disconnected from slave")
 
     def killTask(self, driver, taskId):
         if taskId in self.runningTasks:
             os.kill(self.runningTasks[taskId], 9)
 
     def shutdown(self, driver):
-        log.critical("Shutting Down Executor...")
+        log.critical("Shutting down executor...")
         for taskId, pid in self.runningTasks.items():
             self.killTask(driver, taskId)
-        log.critical("Executor Shut Down")
+        Resource.cleanSystem()
+        log.critical("Executor shut down")
 
     def error(self, driver, message):
         """
@@ -88,8 +89,8 @@ def _sendStats(self, driver):
         while True:
             coresUsage = str(psutil.cpu_percent())
             ramUsage = str(psutil.virtual_memory().percent)
-            driver.sendFrameworkMessage("CPU usage: %s, Memory usage: %s" % (coresUsage, ramUsage))
-            log.debug("sent stats message")
+            driver.sendFrameworkMessage("CPU usage: %s, memory usage: %s" % (coresUsage, ramUsage))
+            log.debug("Sent stats message")
             sleep(30)
 
     def launchTask(self, driver, task):
diff --git a/src/toil/resource.py b/src/toil/resource.py
--- a/src/toil/resource.py
+++ b/src/toil/resource.py
@@ -25,11 +25,12 @@
 import errno
 from zipfile import ZipFile
 import sys
+import shutil
 
-log = logging.getLogger( __name__ )
+log = logging.getLogger(__name__)
 
 
-class Resource( namedtuple( 'Resource', ('name', 'pathHash', 'url', 'contentHash') ) ):
+class Resource(namedtuple('Resource', ('name', 'pathHash', 'url', 'contentHash'))):
     """
     Represents a file or directory that will be deployed to each node before any jobs in the user
     script are invoked. Each instance is a namedtuple with the following elements:
@@ -53,7 +54,7 @@ class Resource( namedtuple( 'Resource', ('name', 'pathHash', 'url', 'contentHash
     resourceEnvNamePrefix = 'JTRES_'
 
     @classmethod
-    def create( cls, jobStore, leaderPath ):
+    def create(cls, jobStore, leaderPath):
         """
         Saves the content of the file or directory at the given path to the given job store
         and returns a resource object representing that content for the purpose of obtaining it
@@ -61,51 +62,63 @@ def create( cls, jobStore, leaderPath ):
 
         :rtype: Resource
         """
-        if os.path.isdir( leaderPath ):
+        if os.path.isdir(leaderPath):
             subcls = DirectoryResource
-        elif os.path.isfile( leaderPath ):
+        elif os.path.isfile(leaderPath):
             subcls = FileResource
-        elif os.path.exists( leaderPath ):
-            raise AssertionError( "Neither a file or a directory: '%s'" % leaderPath )
+        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 )
+            raise AssertionError("No such file or directory: '%s'" % leaderPath)
 
-        pathHash = subcls._pathHash( leaderPath )
-        contentHash = hashlib.md5( )
+        pathHash = subcls._pathHash(leaderPath)
+        contentHash = hashlib.md5()
         # noinspection PyProtectedMember
-        with subcls._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( ) )
+        with subcls._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())
 
     @classmethod
-    def prepareSystem( cls ):
+    def prepareSystem(cls):
         """
         Prepares this system for the downloading and lookup of resources. This method should only
         be invoked on a worker node. It is idempotent but not thread-safe.
         """
         try:
-            resourceRootDirPath = os.environ[ cls.rootDirPathEnvName ]
+            resourceRootDirPath = os.environ[cls.rootDirPathEnvName]
         except KeyError:
             # Create directory holding local copies of requested resources ...
-            resourceRootDirPath = mkdtemp( )
+            resourceRootDirPath = mkdtemp()
             # .. 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 )
+            os.environ[cls.rootDirPathEnvName] = resourceRootDirPath
+        assert os.path.isdir(resourceRootDirPath)
 
-    def register( self ):
+    @classmethod
+    def cleanSystem(cls):
+        """
+        Removes all downloaded, localized resources
+        """
+        resourceRootDirPath = os.environ[cls.rootDirPathEnvName]
+        shutil.rmtree( resourceRootDirPath )
+        for k,v in os.environ.items():
+            if k.startswith( cls.resourceEnvNamePrefix ):
+                os.unsetenv(k)
+
+
+    def register(self):
         """
         Register this resource for later retrieval via lookup(), possibly in a child process.
         """
-        os.environ[ self.resourceEnvNamePrefix + self.pathHash ] = self._pickle( )
+        os.environ[self.resourceEnvNamePrefix + self.pathHash] = self._pickle()
 
     @classmethod
-    def lookup( cls, leaderPath ):
+    def lookup(cls, leaderPath):
         """
         Returns a resource object representing a resource created from a file or directory at the
         given path on the leader. This method should be invoked on the worker. The given path
@@ -115,31 +128,31 @@ def lookup( cls, leaderPath ):
 
         :rtype: Resource
         """
-        pathHash = cls._pathHash( leaderPath )
+        pathHash = cls._pathHash(leaderPath)
         try:
-            s = os.environ[ cls.resourceEnvNamePrefix + pathHash ]
+            s = os.environ[cls.resourceEnvNamePrefix + pathHash]
         except KeyError:
-            log.warn( "Can't find resource for leader path '%s'", leaderPath )
+            log.warn("Can't find resource for leader path '%s'", leaderPath)
             return None
         else:
-            self = cls._unpickle( s )
+            self = cls._unpickle(s)
             assert self.pathHash == pathHash
             return self
 
-    def download( self, callback=None ):
+    def download(self, callback=None):
         """
         Downloads this resource from its URL to a file on the local system. This method should
         only be invoked on a worker node after the node was setup for accessing resources via
         prepareSystem().
         """
         dirPath = self.localDirPath
-        if not os.path.exists( dirPath ):
-            tempDirPath = mkdtemp( dir=os.path.dirname( dirPath ), prefix=self.contentHash + "-" )
-            self._save( tempDirPath )
+        if not os.path.exists(dirPath):
+            tempDirPath = mkdtemp(dir=os.path.dirname(dirPath), prefix=self.contentHash + "-")
+            self._save(tempDirPath)
             if callback is not None:
-                callback( tempDirPath )
+                callback(tempDirPath)
             try:
-                os.rename( tempDirPath, dirPath )
+                os.rename(tempDirPath, dirPath)
             except OSError as e:
                 if e.errno == errno.ENOTEMPTY:
                     # Another process beat us to it.
@@ -149,7 +162,7 @@ def download( self, callback=None ):
                     raise
 
     @property
-    def localPath( self ):
+    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.
@@ -157,30 +170,30 @@ def localPath( self ):
         raise NotImplementedError
 
     @property
-    def localDirPath( self ):
+    def localDirPath(self):
         """
         The path to the directory containing the resource on the worker.
         """
-        rootDirPath = os.environ[ self.rootDirPathEnvName ]
-        return os.path.join( rootDirPath, self.contentHash )
+        rootDirPath = os.environ[self.rootDirPathEnvName]
+        return os.path.join(rootDirPath, self.contentHash)
 
-    def _pickle( self ):
-        return self.__class__.__module__ + "." + self.__class__.__name__ + ':' + json.dumps( self )
+    def _pickle(self):
+        return self.__class__.__module__ + "." + self.__class__.__name__ + ':' + json.dumps(self)
 
     @classmethod
-    def _unpickle( cls, s ):
+    def _unpickle(cls, s):
         """
         :rtype: Resource
         """
-        className, _json = s.split( ':', 1 )
-        return locate( className )( *json.loads( _json ) )
+        className, _json = s.split(':', 1)
+        return locate(className)(*json.loads(_json))
 
     @classmethod
-    def _pathHash( cls, path ):
-        return hashlib.md5( path ).hexdigest( )
+    def _pathHash(cls, path):
+        return hashlib.md5(path).hexdigest()
 
     @classmethod
-    def _load( cls, path ):
+    def _load(cls, path):
         """
         Returns a readable file-like object for the given path. If the path refers to a regular
         file, this method returns the result of invoking open() on the given path. If the path
@@ -190,81 +203,81 @@ def _load( cls, path ):
         :type path: str
         :rtype: io.IOBase
         """
-        raise NotImplementedError( )
+        raise NotImplementedError()
 
-    def _save( self, dirPath ):
+    def _save(self, dirPath):
         """
         Save this resource to disk at the given parent path.
 
         :type dirPath: str
         """
-        raise NotImplementedError( )
+        raise NotImplementedError()
 
-    def _download( self, dstFile ):
+    def _download(self, dstFile):
         """
         Download this resource from its URL to the given file object.
 
         :type dstFile: io.BytesIO|io.FileIO
         """
-        with closing( urlopen( self.url ) ) as content:
-            buf = content.read( )
-        contentHash = hashlib.md5( buf )
-        assert contentHash.hexdigest( ) == self.contentHash
-        dstFile.write( buf )
+        with closing(urlopen(self.url)) as content:
+            buf = content.read()
+        contentHash = hashlib.md5(buf)
+        assert contentHash.hexdigest() == self.contentHash
+        dstFile.write(buf)
 
 
-class FileResource( Resource ):
+class FileResource(Resource):
     """
     A resource read from a file on the leader.
     """
 
     @classmethod
-    def _load( cls, path ):
-        return open( path )
+    def _load(cls, path):
+        return open(path)
 
-    def _save( self, dirPath ):
-        with open( os.path.join( dirPath, self.name ), mode='w' ) as localFile:
-            self._download( localFile )
+    def _save(self, dirPath):
+        with open(os.path.join(dirPath, self.name), mode='w') as localFile:
+            self._download(localFile)
 
     @property
-    def localPath( self ):
-        return os.path.join( self.localDirPath, self.name )
+    def localPath(self):
+        return os.path.join(self.localDirPath, self.name)
 
 
-class DirectoryResource( Resource ):
+class DirectoryResource(Resource):
     """
     A resource read from a directory on the leader. The URL will point to a ZIP archive of the
     directory.
     """
 
     @classmethod
-    def _load( cls, leaderPath ):
-        bytesIO = BytesIO( )
-        with ZipFile( file=bytesIO, mode='w' ) as zipFile:
-            for dirPath, fileNames, dirNames in os.walk( leaderPath ):
-                assert dirPath.startswith( leaderPath )
+    def _load(cls, leaderPath):
+        bytesIO = BytesIO()
+        with ZipFile(file=bytesIO, mode='w') as zipFile:
+            for dirPath, fileNames, dirNames in os.walk(leaderPath):
+                assert dirPath.startswith(leaderPath)
                 for fileName in fileNames:
-                    filePath = os.path.join( dirPath, fileName )
-                    assert filePath.encode( 'ascii' ) == filePath
-                    relativeFilePath = os.path.relpath( filePath, leaderPath )
-                    assert not relativeFilePath.startswith( os.path.sep )
-                    zipFile.write( filePath, relativeFilePath )
-        bytesIO.seek( 0 )
+                    filePath = os.path.join(dirPath, fileName)
+                    assert filePath.encode('ascii') == filePath
+                    relativeFilePath = os.path.relpath(filePath, leaderPath)
+                    assert not relativeFilePath.startswith(os.path.sep)
+                    zipFile.write(filePath, relativeFilePath)
+        bytesIO.seek(0)
         return bytesIO
 
-    def _save( self, dirPath ):
-        bytesIO = BytesIO( )
-        self._download( bytesIO )
-        bytesIO.seek( 0 )
-        with ZipFile( file=bytesIO, mode='r' ) as zipFile:
-            zipFile.extractall( path=dirPath )
+    def _save(self, dirPath):
+        bytesIO = BytesIO()
+        self._download(bytesIO)
+        bytesIO.seek(0)
+        with ZipFile(file=bytesIO, mode='r') as zipFile:
+            zipFile.extractall(path=dirPath)
 
     @property
-    def localPath( self ):
+    def localPath(self):
         return self.localDirPath
 
 
-class ModuleDescriptor( namedtuple( 'ModuleDescriptor', ('dirPath', 'name', 'extension') ) ):
+class ModuleDescriptor(namedtuple('ModuleDescriptor', ('dirPath', 'name', 'extension'))):
     """
     A path to a Python module decomposed into a namedtuple of three elements, namely
 
@@ -293,41 +306,41 @@ class ModuleDescriptor( namedtuple( 'ModuleDescriptor', ('dirPath', 'name', 'ext
     """
 
     @classmethod
-    def forModule( cls, moduleName ):
+    def forModule(cls, moduleName):
         """
         Return an instance of this class representing the module of the given name. If the given
         module name is "__main__", it will be translated to the actual file name of the top-level
         script without the .py or .pyc extension. This method assumes that the module with the
         specified name has already been loaded.
         """
-        module = sys.modules[ moduleName ]
-        moduleFilePath = os.path.abspath( module.__file__ )
-        dirPath = moduleFilePath.split( os.path.sep )
-        dirPath[ -1 ], extension = os.path.splitext( dirPath[ -1 ] )
+        module = sys.modules[moduleName]
+        moduleFilePath = os.path.abspath(module.__file__)
+        dirPath = moduleFilePath.split(os.path.sep)
+        dirPath[-1], extension = os.path.splitext(dirPath[-1])
         assert extension in ('.py', '.pyc')
         if moduleName == '__main__':
-            moduleName = dirPath.pop( )
+            moduleName = dirPath.pop()
         else:
-            for package in reversed( moduleName.split( '.' ) ):
-                dirPathTail = dirPath.pop( )
+            for package in reversed(moduleName.split('.')):
+                dirPathTail = dirPath.pop()
                 assert dirPathTail == package
-        return cls( dirPath=os.path.sep.join( dirPath ), name=moduleName, extension=extension )
+        return cls(dirPath=os.path.sep.join(dirPath), name=moduleName, extension=extension)
 
     @property
-    def belongsToToil( self ):
+    def belongsToToil(self):
         """
         True if this module is part of the Toil distribution
         """
-        return self.name.startswith( 'toil.' )
+        return self.name.startswith('toil.')
 
     @property
-    def filePath( self ):
+    def filePath(self):
         """
         The full path to the file containing this module
         """
-        return os.path.join( self.dirPath, *self.name.split( '.' ) ) + self.extension
+        return os.path.join(self.dirPath, *self.name.split('.')) + self.extension
 
-    def saveAsResourceTo( self, jobStore ):
+    def saveAsResourceTo(self, jobStore):
         """
         Store the file containing this module--or even the Python package directory hierarchy
         containing that file--as a resource to the given job store and return the
@@ -336,42 +349,42 @@ def saveAsResourceTo( self, jobStore ):
         :type jobStore: toil.jobStores.abstractJobStore.AbstractJobStore
         :rtype: toil.resource.Resource
         """
-        return Resource.create( jobStore, self._resourcePath )
+        return Resource.create(jobStore, self._resourcePath)
 
-    def localize( self ):
+    def localize(self):
         """
         Check if this module was saved as a resource. If it was, return a new module descriptor
         that points to a local copy of that resource. Should only be called on a worker node.
 
         :rtype: toil.resource.Resource
         """
-        resource = Resource.lookup( self._resourcePath )
+        resource = Resource.lookup(self._resourcePath)
         if resource is None:
-            log.warn( "Can't localize module %r", self )
+            log.warn("Can't localize module %r", self)
             return self
         else:
-            def stash( tmpDirPath ):
+            def stash(tmpDirPath):
                 # Save the original dirPath such that we can restore it in globalize()
-                with open( os.path.join( tmpDirPath, '.original' ), 'w' ) as f:
-                    f.write( json.dumps( self ) )
+                with open(os.path.join(tmpDirPath, '.original'), 'w') as f:
+                    f.write(json.dumps(self))
 
-            resource.download( callback=stash )
-            return self.__class__( dirPath=resource.localDirPath, name=self.name,
-                                   extension=self.extension )
+            resource.download(callback=stash)
+            return self.__class__(dirPath=resource.localDirPath, name=self.name,
+                                  extension=self.extension)
 
-    def globalize( self ):
+    def globalize(self):
         try:
-            with open( os.path.join( self.dirPath, '.original' ) ) as f:
-                return self.__class__( *json.loads( f.read( ) ) )
+            with open(os.path.join(self.dirPath, '.original')) as f:
+                return self.__class__(*json.loads(f.read()))
         except IOError as e:
             if e.errno == errno.ENOENT:
-                log.warn( "Can't globalize module %r.", self )
+                log.warn("Can't globalize module %r.", self)
                 return self
             else:
                 raise
 
     @property
-    def _resourcePath( self ):
+    def _resourcePath(self):
         """
         The path to the file or package directory that should be used when shipping this module
         around as a resource.


Test patch: