Last active
October 9, 2015 13:28
-
-
Save tubaman/3515397 to your computer and use it in GitHub Desktop.
Disco worker that zips up imported modules/packages
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
I've updated the gist with an example in the Worker class docstring