Path for run_sorters with docker images
@@ -16,13 +16,12 @@ from .sorterlist import sorter_dict
def _run_one(arg_list):
# the multiprocessing python module force to have one unique tuple argument
- sorter_name, recording, output_folder, verbose, sorter_params = arg_list
+ sorter_name, recording, output_folder, verbose, sorter_params, docker_image = arg_list
if isinstance(recording, dict):
recording = load_extractor(recording)
else:
recording = recording
- SorterClass = sorter_dict[sorter_name]
# because this is checks in run_sorters before this call
remove_existing_folder = False
@@ -31,11 +30,16 @@ def _run_one(arg_list):
# because we won't want the loop/worker to break
raise_error = False
- # only classmethod call not instance (stateless at instance level but state is in folder)
- output_folder = SorterClass.initialize_folder(recording, output_folder, verbose, remove_existing_folder)
- SorterClass.set_params_to_folder(recording, output_folder, sorter_params, verbose)
- SorterClass.setup_recording(recording, output_folder, verbose=verbose)
- SorterClass.run_from_folder(output_folder, raise_error, verbose)
+ if docker_image is None:
+
+ run_sorter_local(sorter_name, recording, output_folder, remove_existing_folder,
+ delete_output_folder, verbose, raise_error)
+
+ else:
+
+ run_sorter_docker(sorter_name, recording, docker_image, output_folder=output_folder,
+ remove_existing_folder=remove_existing_folder, delete_output_folder=delete_output_folder,
+ verbose=verbose, raise_error=raise_error, **sorter_params):
_implemented_engine = ('loop', 'joblib', 'dask')
@@ -50,6 +54,7 @@ def run_sorters(sorter_list,
engine_kwargs={},
verbose=False,
with_output=True,
+ docker_images = {}
):
"""
This run several sorter on several recording.
@@ -161,13 +166,20 @@ def run_sorters(sorter_list,
shutil.rmtree(str(output_folder))
params = sorter_params.get(sorter_name, {})
+ use_docker = sorter_name in docker_images
+ if use_docker:
+ docker_image = docker_images[sorter_name]
+ else:
+ docker_image = None
+
if need_dump:
if not recording.is_dumpable:
raise Exception('recording not dumpable call recording.save() before')
recording_arg = recording.to_dict()
else:
recording_arg = recording
- task_args = (sorter_name, recording_arg, output_folder, verbose, params)
+
+ task_args = (sorter_name, recording_arg, output_folder, verbose, params, docker_image)
task_args_list.append(task_args)
if engine == 'loop':
