Disco worker that zips up imported modules/packages

worker.py

#!/usr/bin/env python
"""Disco worker that automatically finds and sends to the nodes all the modules and packages that are used"""
import os
import sys
import shutil
import tempfile
import zipfile
import modulefinder
from inspect import getsourcefile, getmodule

import disco.worker.classic.worker

def should_ignore(path):
    path_prefixes_to_ignore = ['/usr', '/opt']
    for prefix in path_prefixes_to_ignore:
        if path is not None and path.startswith(prefix):
            return True
    return False


def make_modules_zipfile(filename):
    zip_dir = tempfile.mkdtemp()
    zip_filename = os.path.join(zip_dir, 'modules.zip')
    finder = modulefinder.ModuleFinder()
    finder.run_script(filename)
    zf = zipfile.ZipFile(zip_filename, 'w')
    for name, mod in finder.modules.iteritems():
        if not mod.__file__:
            continue
        modfile = os.path.abspath(mod.__file__)
        if not should_ignore(modfile):
            for path in sys.path:
                if os.path.commonprefix([path, modfile]) == path:
                    relative_path = os.path.relpath(modfile, path)
                    break
            zf.write(modfile, relative_path)
    zf.close()
    return zip_dir, zip_filename


class Worker(disco.worker.classic.worker.Worker):
   """ex:
   class MyJob(Job):
       from worker import Worker
       ...
    
    """
    def jobenvs(self, job, **jobargs):
        envs = super(Worker, self).jobenvs(job, **jobargs)
        envs['PYTHONPATH'] = ':'.join(('lib/modules.zip',
            envs.get('PYTHONPATH', '')))
        return envs

    def jobzip(self, job, **jobargs):
        def get(key):
            return self.getitem(key, job, jobargs)
        required_files = get('required_files')
        job_path = getsourcefile(getmodule(job))
        zip_dir, modules_zipfile = make_modules_zipfile(job_path)
        try:
            if isinstance(required_files, dict):
                zipdata = open(modules_zipfile, 'rb').read()
                required_files['lib/modules.zip'] = zipdata
            else:
                required_files.append(modules_zipfile)
            return super(Worker, self).jobzip(job, **jobargs)
        finally:
            if os.path.exists(zip_dir):
                shutil.rmtree(zip_dir)


if __name__ == '__main__':
    Worker.main()

This worker collects all required modules(including packages), zips them, up and sends them to the nodes. It packages up everything not under '/usr' or '/opt'. Any modules under /usr or /opt are assumed to be installed on every node.

Can you give an example of using this worker? I've tried it but it's not detecting the modules I need it to detect.

I've updated the gist with an example in the Worker class docstring