Repository: SpikeInterface/spikeinterface (python)

Issue:
Sorting analyzer extensions don't know if they've been interrupted
Hi all — sometimes a sorting analyzer's computation of an extension is interrupted, either accidentally (i.e. a SLURM job times out), or manually (I want to change a parameter). Normally, the way one would deal with this, in a script, is to check if the extension has been completed with `get_extension`/`has_extension`, but these only check for the presence of the corresponding dir and the file `params.json` (see [here](https://github.com/SpikeInterface/spikeinterface/blob/39218069ca61529575abf70095d1ce7912a410d3/src/spikeinterface/core/sortinganalyzer.py#L1364)), and so if the computation has been interrupted partway through, the code would treat it as completed, but downstream analyses would fail / be wrong. Currently, for me, I'm manually deleting extension folders if this occurs and I need to fix it. Is there an existing better solution? Otherwise, I'd propose simply having a `is_finished` or `exit_status` flag in `info.json` that could be used for `has_extension`, rather than just checking that the folder / params are there. Thanks!

Pull request description:
Analyzer extension exit status
Fixes #3329

Here's a first pass at adding information about an extension's completion (or not) to its metadata.

This PR adds a `run_info.json` file to each extension's folder. This contains keys:
* `run_completed` (whether or not the call to `AnalyzerExtension._run()` finished),
* `data_loadable` (whether or not `AnalyzerExtension.load_data()` can be called without an error) 
* `runtime_s` (the runtime of `AnalyzerExtension._run()` in seconds). 

The core logic is implemented in `AnalyzerExtension.run()` and wraps the call to the extension-specific `_run()`. Then when `AnalyzerExtension.get_extension()` is called (and redirects to `AnalyzerExtension.load()` internally), if the run wasn't completed, it simply returns None as if the extension doesn't exist, and the user would be able to catch that and re-run the extension.
 
There is a mild complication with extensions that use the `run_node_pipeline` if called through `compute_several_extensions`, because then `AnalyzerExtension.run()` is never actually called. Instead, it looks like `AnalyzerExtension.save()` is used, so I added the relevant lines there to catch that, and assume that if the code makes it there, the run is completed.

TODO:
* I'm not sure if the implementations in `merge()` and `copy()` are correct, I don't totally understand when those methods are expected to be called / what's wrapping them.
* Write some tests?
 
Notes:
* It is a bit of a slow-down to check if the data are loadable after each run. I'm ambivalent as to whether that check is actually useful, especially since data corruption is more likely to happen sometimes significantly after the run, not immediately after.
* It seemed more reasonable to add a separate `run_info.json` file than trying to cram this into `info.json`, since that is more about the code itself.


Patch:
diff --git a/src/spikeinterface/core/sortinganalyzer.py b/src/spikeinterface/core/sortinganalyzer.py
index 13b64664a..57c9e5f37 100644
--- a/src/spikeinterface/core/sortinganalyzer.py
+++ b/src/spikeinterface/core/sortinganalyzer.py
@@ -11,6 +11,7 @@ import weakref
 import shutil
 import warnings
 import importlib
+from time import perf_counter
 
 import numpy as np
 
@@ -1336,6 +1337,7 @@ extension_params={"waveforms":{"ms_before":1.5, "ms_after": "2.5"}}\
 
             job_name = "Compute : " + " + ".join(extensions_with_pipeline.keys())
 
+            t_start = perf_counter()
             results = run_node_pipeline(
                 self.recording,
                 all_nodes,
@@ -1345,10 +1347,15 @@ extension_params={"waveforms":{"ms_before":1.5, "ms_after": "2.5"}}\
                 squeeze_output=False,
                 verbose=verbose,
             )
+            t_end = perf_counter()
+            # for pipeline node extensions we can only track the runtime of the run_node_pipeline
+            runtime_s = t_end - t_start
 
             for r, result in enumerate(results):
                 extension_name, variable_name = result_routage[r]
                 extension_instances[extension_name].data[variable_name] = result
+                extension_instances[extension_name].run_info["runtime_s"] = runtime_s
+                extension_instances[extension_name].run_info["run_completed"] = True
 
             for extension_name, extension_instance in extension_instances.items():
                 self.extensions[extension_name] = extension_instance
@@ -1738,8 +1745,12 @@ class AnalyzerExtension:
         self._sorting_analyzer = weakref.ref(sorting_analyzer)
 
         self.params = None
+        self.run_info = self._default_run_info_dict()
         self.data = dict()
 
+    def _default_run_info_dict(self):
+        return dict(run_completed=False, runtime_s=None)
+
     #######
     # This 3 methods must be implemented in the subclass!!!
     # See DummyAnalyzerExtension in test_sortinganalyzer.py as a simple example
@@ -1851,11 +1862,42 @@ class AnalyzerExtension:
     def load(cls, sorting_analyzer):
         ext = cls(sorting_analyzer)
         ext.load_params()
-        ext.load_data()
-        if cls.need_backward_compatibility_on_load:
-            ext._handle_backward_compatibility_on_load()
+        ext.load_run_info()
+        if ext.run_info is not None:
+            if ext.run_info["run_completed"]:
+                ext.load_data()
+                if cls.need_backward_compatibility_on_load:
+                    ext._handle_backward_compatibility_on_load()
+                if len(ext.data) > 0:
+                    return ext
+        else:
+            # this is for back-compatibility of old analyzers
+            ext.load_data()
+            if cls.need_backward_compatibility_on_load:
+                ext._handle_backward_compatibility_on_load()
+            if len(ext.data) > 0:
+                return ext
+        # If extension run not completed, or data has gone missing,
+        # return None to indicate that the extension should be (re)computed.
+        return None
+
+    def load_run_info(self):
+        if self.format == "binary_folder":
+            extension_folder = self._get_binary_extension_folder()
+            run_info_file = extension_folder / "run_info.json"
+            if run_info_file.is_file():
+                with open(str(run_info_file), "r") as f:
+                    run_info = json.load(f)
+            else:
+                warnings.warn(f"Found no run_info file for {self.extension_name}, extension should be re-computed.")
+                run_info = None
 
-        return ext
+        elif self.format == "zarr":
+            extension_group = self._get_zarr_extension_group(mode="r")
+            run_info = extension_group.attrs.get("run_info", None)
+        if run_info is None:
+            warnings.warn(f"Found no run_info file for {self.extension_name}, extension should be re-computed.")
+        self.run_info = run_info
 
     def load_params(self):
         if self.format == "binary_folder":
@@ -1873,12 +1915,17 @@ class AnalyzerExtension:
         self.params = params
 
     def load_data(self):
+        ext_data = None
         if self.format == "binary_folder":
             extension_folder = self._get_binary_extension_folder()
             for ext_data_file in extension_folder.iterdir():
                 # patch for https://github.com/SpikeInterface/spikeinterface/issues/3041
                 # maybe add a check for version number from the info.json during loading only
-                if ext_data_file.name == "params.json" or ext_data_file.name == "info.json":
+                if (
+                    ext_data_file.name == "params.json"
+                    or ext_data_file.name == "info.json"
+                    or ext_data_file.name == "run_info.json"
+                ):
                     continue
                 ext_data_name = ext_data_file.stem
                 if ext_data_file.suffix == ".json":
@@ -1919,6 +1966,9 @@ class AnalyzerExtension:
                     ext_data = np.array(ext_data_)
                 self.data[ext_data_name] = ext_data
 
+        if len(self.data) == 0:
+            warnings.warn(f"Found no data for {self.extension_name}, extension should be re-computed.")
+
     def copy(self, new_sorting_analyzer, unit_ids=None):
         # alessio : please note that this also replace the old select_units!!!
         new_extension = self.__class__(new_sorting_analyzer)
@@ -1927,6 +1977,7 @@ class AnalyzerExtension:
             new_extension.data = self.data
         else:
             new_extension.data = self._select_extension_data(unit_ids)
+        new_extension.run_info = self.run_info.copy()
         new_extension.save()
         return new_extension
 
@@ -1944,24 +1995,33 @@ class AnalyzerExtension:
         new_extension.data = self._merge_extension_data(
             merge_unit_groups, new_unit_ids, new_sorting_analyzer, keep_mask, verbose=verbose, **job_kwargs
         )
+        new_extension.run_info = self.run_info.copy()
         new_extension.save()
         return new_extension
 
     def run(self, save=True, **kwargs):
         if save and not self.sorting_analyzer.is_read_only():
-            # this also reset the folder or zarr group
+            # NB: this call to _save_params() also resets the folder or zarr group
             self._save_params()
             self._save_importing_provenance()
+            self._save_run_info()
 
+        t_start = perf_counter()
         self._run(**kwargs)
+        t_end = perf_counter()
+        self.run_info["runtime_s"] = t_end - t_start
 
         if save and not self.sorting_analyzer.is_read_only():
             self._save_data(**kwargs)
 
+        self.run_info["run_completed"] = True
+        self._save_run_info()
+
     def save(self, **kwargs):
         self._save_params()
         self._save_importing_provenance()
         self._save_data(**kwargs)
+        self._save_run_info()
 
     def _save_data(self, **kwargs):
         if self.format == "memory":
@@ -2060,6 +2120,7 @@ class AnalyzerExtension:
         """
         self._reset_extension_folder()
         self.params = None
+        self.run_info = self._default_run_info_dict()
         self.data = dict()
 
     def set_params(self, save=True, **params):
@@ -2080,6 +2141,7 @@ class AnalyzerExtension:
         if save:
             self._save_params()
             self._save_importing_provenance()
+            self._save_run_info()
 
     def _save_params(self):
         params_to_save = self.params.copy()
@@ -2117,6 +2179,17 @@ class AnalyzerExtension:
             extension_group = self._get_zarr_extension_group(mode="r+")
             extension_group.attrs["info"] = info
 
+    def _save_run_info(self):
+        run_info = self.run_info.copy()
+
+        if self.format == "binary_folder":
+            extension_folder = self._get_binary_extension_folder()
+            run_info_file = extension_folder / "run_info.json"
+            run_info_file.write_text(json.dumps(run_info, indent=4), encoding="utf8")
+        elif self.format == "zarr":
+            extension_group = self._get_zarr_extension_group(mode="r+")
+            extension_group.attrs["run_info"] = run_info
+
     def get_pipeline_nodes(self):
         assert (
             self.use_nodepipeline
@@ -2124,7 +2197,10 @@ class AnalyzerExtension:
         return self._get_pipeline_nodes()
 
     def get_data(self, *args, **kwargs):
-        assert len(self.data) > 0, f"You must run the extension {self.extension_name} before retrieving data"
+        assert self.run_info[
+            "run_completed"
+        ], f"You must run the extension {self.extension_name} before retrieving data"
+        assert len(self.data) > 0, "Extension has been run but no data found."
         return self._get_data(*args, **kwargs)
 
 


Test patch:
diff --git a/src/spikeinterface/core/tests/test_sortinganalyzer.py b/src/spikeinterface/core/tests/test_sortinganalyzer.py
index 689073d6b..3f45487f4 100644
--- a/src/spikeinterface/core/tests/test_sortinganalyzer.py
+++ b/src/spikeinterface/core/tests/test_sortinganalyzer.py
@@ -125,6 +125,39 @@ def test_SortingAnalyzer_zarr(tmp_path, dataset):
     )
 
 
+def test_load_without_runtime_info(tmp_path, dataset):
+    recording, sorting = dataset
+
+    folder = tmp_path / "test_SortingAnalyzer_run_info"
+
+    extensions = ["random_spikes", "templates"]
+    # binary_folder
+    sorting_analyzer = create_sorting_analyzer(
+        sorting, recording, format="binary_folder", folder=folder, sparse=False, sparsity=None
+    )
+    sorting_analyzer.compute(extensions)
+    # remove run_info.json to mimic a previous version of spikeinterface
+    for ext in extensions:
+        (folder / "extensions" / ext / "run_info.json").unlink()
+    # should raise a warning for missing run_info
+    with pytest.warns(UserWarning):
+        sorting_analyzer = load_sorting_analyzer(folder, format="auto")
+
+    # zarr
+    folder = tmp_path / "test_SortingAnalyzer_run_info.zarr"
+    sorting_analyzer = create_sorting_analyzer(
+        sorting, recording, format="zarr", folder=folder, sparse=False, sparsity=None
+    )
+    sorting_analyzer.compute(extensions)
+    # remove run_info from attrs to mimic a previous version of spikeinterface
+    root = sorting_analyzer._get_zarr_root(mode="r+")
+    for ext in extensions:
+        del root["extensions"][ext].attrs["run_info"]
+    # should raise a warning for missing run_info
+    with pytest.warns(UserWarning):
+        sorting_analyzer = load_sorting_analyzer(folder, format="auto")
+
+
 def test_SortingAnalyzer_tmp_recording(dataset):
     recording, sorting = dataset
     recording_cached = recording.save(mode="memory")