From fdbd852b38bf932eff016a8c5e0bddbecd7c5171 Mon Sep 17 00:00:00 2001 From: Samuel Garcia Date: Wed, 29 Nov 2023 12:04:08 +0100 Subject: [PATCH 1/2] run_sorter in container check json or pickle --- src/spikeinterface/sorters/runsorter.py | 26 ++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/src/spikeinterface/sorters/runsorter.py b/src/spikeinterface/sorters/runsorter.py index da8b48085c..0a13f8d754 100644 --- a/src/spikeinterface/sorters/runsorter.py +++ b/src/spikeinterface/sorters/runsorter.py @@ -2,6 +2,7 @@ import os from pathlib import Path import json +import pickle import platform from warnings import warn from typing import Optional, Union @@ -414,9 +415,15 @@ def run_sorter_container( # create 3 files for communication with container # recording dict inside - (parent_folder / "in_container_recording.json").write_text( - json.dumps(check_json(rec_dict), indent=4), encoding="utf8" - ) + if recording.check_serializability("json"): + (parent_folder / "in_container_recording.json").write_text( + json.dumps(check_json(rec_dict), indent=4), encoding="utf8" + ) + elif recording.check_serializability("pickle"): + (parent_folder / "in_container_recording.pickle").write_bytes(pickle.dumps(rec_dict)) + else: + raise RuntimeError("To use run_sorter with container the recording must serializable") + # need to share specific parameters (parent_folder / "in_container_params.json").write_text( json.dumps(check_json(sorter_params), indent=4), encoding="utf8" @@ -433,13 +440,19 @@ def run_sorter_container( # the py script py_script = f""" import json +from pathlib import Path from spikeinterface import load_extractor from spikeinterface.sorters import run_sorter_local if __name__ == '__main__': # this __name__ protection help in some case with multiprocessing (for instance HS2) # load recording in container - recording = load_extractor('{parent_folder_unix}/in_container_recording.json') + json_rec = Path('{parent_folder_unix}/in_container_recording.json') + pickle_rec = Path('{parent_folder_unix}/in_container_recording.pickle') + if json_rec.exists(): + recording = load_extractor(json_rec) + else: + recording = load_extractor(pickle_rec) # load params in container with open('{parent_folder_unix}/in_container_params.json', encoding='utf8', mode='r') as f: @@ -593,7 +606,10 @@ def run_sorter_container( # clean useless files if delete_container_files: - os.remove(parent_folder / "in_container_recording.json") + if (parent_folder / "in_container_recording.json").exists(): + os.remove(parent_folder / "in_container_recording.json") + if (parent_folder / "in_container_recording.pickle").exists(): + os.remove(parent_folder / "in_container_recording.pickle") os.remove(parent_folder / "in_container_params.json") os.remove(parent_folder / "in_container_sorter_script.py") if mode == "singularity": From 2f7ca19b7cc15ec72ba058b30388f6fc1a0378b3 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Thu, 30 Nov 2023 10:26:44 +0100 Subject: [PATCH 2/2] Update src/spikeinterface/sorters/runsorter.py Co-authored-by: Heberto Mayorquin --- src/spikeinterface/sorters/runsorter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/spikeinterface/sorters/runsorter.py b/src/spikeinterface/sorters/runsorter.py index 0a13f8d754..7591c9eb2c 100644 --- a/src/spikeinterface/sorters/runsorter.py +++ b/src/spikeinterface/sorters/runsorter.py @@ -422,7 +422,7 @@ def run_sorter_container( elif recording.check_serializability("pickle"): (parent_folder / "in_container_recording.pickle").write_bytes(pickle.dumps(rec_dict)) else: - raise RuntimeError("To use run_sorter with container the recording must serializable") + raise RuntimeError("To use run_sorter with container the recording must be serializable") # need to share specific parameters (parent_folder / "in_container_params.json").write_text(