Skip to content
Snippets Groups Projects
Commit e90dec37 authored by Michiel Cottaar's avatar Michiel Cottaar
Browse files

Use dill to store function on disc

Dill is a lot more robust than my previous attempt to retreive
functions stored on disc by func_to_cmd.
parent 9a0d79c5
No related branches found
No related tags found
No related merge requests found
...@@ -41,7 +41,7 @@ from io import BytesIO ...@@ -41,7 +41,7 @@ from io import BytesIO
import os.path as op import os.path as op
import glob import glob
import time import time
import pickle import dill
import sys import sys
import tempfile import tempfile
import logging import logging
...@@ -455,23 +455,12 @@ def hold(job_ids, hold_filename=None): ...@@ -455,23 +455,12 @@ def hold(job_ids, hold_filename=None):
_external_job = ("""#!{} _external_job = ("""#!{}
# This is a temporary file designed to run the python function {}, # This is a temporary file designed to run the python function {},
# so that it can be submitted to the cluster # so that it can be submitted to the cluster
import pickle import dill
from io import BytesIO from io import BytesIO
from importlib import import_module from importlib import import_module
{} {}
pickle_bytes = BytesIO({}) dill_bytes = BytesIO({})
name_type, name, func_name, args, kwargs = pickle.load(pickle_bytes) func, args, kwargs = dill.load(dill_bytes)
if name_type == 'module':
# retrieves a function defined in an external module
func = getattr(import_module(name), func_name)
elif name_type == 'script':
# retrieves a function defined in the __main__ script
local_execute = {{'__name__': '__not_main__', '__file__': name}}
exec(open(name, 'r').read(), local_execute)
func = local_execute[func_name]
else:
raise ValueError('Unknown name_type: %r' % name_type)
{} {}
...@@ -504,13 +493,8 @@ def func_to_cmd(func, args=None, kwargs=None, tmp_dir=None, clean="never", verbo ...@@ -504,13 +493,8 @@ def func_to_cmd(func, args=None, kwargs=None, tmp_dir=None, clean="never", verbo
args = () args = ()
if kwargs is None: if kwargs is None:
kwargs = {} kwargs = {}
pickle_bytes = BytesIO() dill_bytes = BytesIO()
if func.__module__ == '__main__': dill.dump((func, args, kwargs), dill_bytes)
pickle.dump(('script', importlib.import_module('__main__').__file__, func.__name__,
args, kwargs), pickle_bytes)
else:
pickle.dump(('module', func.__module__, func.__name__,
args, kwargs), pickle_bytes)
handle, filename = tempfile.mkstemp(prefix=func.__name__ + '_', handle, filename = tempfile.mkstemp(prefix=func.__name__ + '_',
suffix='.py', suffix='.py',
...@@ -531,7 +515,7 @@ import os; os.remove("{filename}")""" ...@@ -531,7 +515,7 @@ import os; os.remove("{filename}")"""
python_cmd = _external_job.format(sys.executable, python_cmd = _external_job.format(sys.executable,
func.__name__, func.__name__,
verbose_script, verbose_script,
pickle_bytes.getvalue(), dill_bytes.getvalue(),
run_script) run_script)
with open(filename, 'w') as python_file: with open(filename, 'w') as python_file:
......
...@@ -2,3 +2,4 @@ h5py>=2.9 ...@@ -2,3 +2,4 @@ h5py>=2.9
nibabel>=2.4 nibabel>=2.4
numpy>=1 numpy>=1
scipy>=0.18 scipy>=0.18
dill
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment