Use file lock when writing results to prevent overwriting

environment.yml

@@ -8,3 +8,4 @@ dependencies:
   - jupyterlab
   - tqdm
   - psutil
+  - filelock

environment_rtd.yml

@@ -8,3 +8,4 @@ dependencies:
   - tqdm
   - psutil
   - pydata-sphinx-theme < 0.10.0
+  - filelock

mesmerize_core/algorithms/cnmf.py

@@ -12,6 +12,7 @@
 import os
 import time
 from datetime import datetime
+from filelock import FileLock, Timeout
 
 # prevent circular import
 if __name__ in ["__main__", "__mp_main__"]:  # when running in subprocess
@@ -138,13 +139,29 @@ def run_algo(batch_path, uuid, data_path: str = None):
 
     cm.stop_server(dview=dview)
 
-    # Add dictionary to output column of series
-    df.loc[df["uuid"] == uuid, "outputs"] = [d]
-    # Add ran timestamp to ran_time column of series
-    df.loc[df["uuid"] == uuid, "ran_time"] = datetime.now().isoformat(timespec="seconds", sep="T")
-    df.loc[df["uuid"] == uuid, "algo_duration"] = str(round(time.time() - algo_start, 2)) + " sec"
-    # save dataframe to disc
-    df.to_pickle(batch_path)
+    # lock batch file while writing back results
+    batch_lock = FileLock(batch_path + '.lock', timeout=30)
+    try:
+        with batch_lock:
+            df = load_batch(batch_path)
+
+            # Add dictionary to output column of series
+            df.loc[df["uuid"] == uuid, "outputs"] = [d]
+            # Add ran timestamp to ran_time column of series
+            df.loc[df["uuid"] == uuid, "ran_time"] = datetime.now().isoformat(timespec="seconds", sep="T")
+            df.loc[df["uuid"] == uuid, "algo_duration"] = str(round(time.time() - algo_start, 2)) + " sec"
+            # save dataframe to disc
+            df.to_pickle(batch_path)
+    except Timeout:
+        # Print a message with details in lieu of writing to the batch file
+        msg = f"Batch file could not be written to within {batch_lock.timeout} seconds."
+        if d["success"]:
+            msg += f"\nRun succeeded; results are in {output_dir}."
+        else:
+            msg += f"Run failed. Traceback:\n"
+            msg += d["traceback"]
+
+        raise RuntimeError(msg)
 
 
 @click.command()

mesmerize_core/algorithms/cnmfe.py

@@ -10,6 +10,7 @@
 import os
 import time
 from datetime import datetime
+from filelock import FileLock, Timeout
 
 if __name__ in ["__main__", "__mp_main__"]:  # when running in subprocess
     from mesmerize_core import set_parent_raw_data_path, load_batch
@@ -121,13 +122,29 @@ def run_algo(batch_path, uuid, data_path: str = None):
 
     cm.stop_server(dview=dview)
 
-    # Add dictionary to output column of series
-    df.loc[df["uuid"] == uuid, "outputs"] = [d]
-    # Add ran timestamp to ran_time column of series
-    df.loc[df["uuid"] == uuid, "ran_time"] = datetime.now().isoformat(timespec="seconds", sep="T")
-    df.loc[df["uuid"] == uuid, "algo_duration"] = str(round(time.time() - algo_start, 2)) + " sec"
-    # save dataframe to disc
-    df.to_pickle(batch_path)
+    # lock batch file while writing back results
+    batch_lock = FileLock(batch_path + '.lock', timeout=30)
+    try:
+        with batch_lock:
+            df = load_batch(batch_path)
+
+            # Add dictionary to output column of series
+            df.loc[df["uuid"] == uuid, "outputs"] = [d]
+            # Add ran timestamp to ran_time column of series
+            df.loc[df["uuid"] == uuid, "ran_time"] = datetime.now().isoformat(timespec="seconds", sep="T")
+            df.loc[df["uuid"] == uuid, "algo_duration"] = str(round(time.time() - algo_start, 2)) + " sec"
+            # save dataframe to disc
+            df.to_pickle(batch_path)
+    except Timeout:
+        # Print a message with details in lieu of writing to the batch file
+        msg = f"Batch file could not be written to within {batch_lock.timeout} seconds."
+        if d["success"]:
+            msg += f"\nRun succeeded; results are in {output_dir}."
+        else:
+            msg += f"Run failed. Traceback:\n"
+            msg += d["traceback"]
+
+        raise RuntimeError(msg)
 
 
 @click.command()

mesmerize_core/algorithms/mcorr.py

@@ -12,6 +12,7 @@
 from shutil import move as move_file
 import time
 from datetime import datetime
+from filelock import FileLock, Timeout
 
 
 # prevent circular import
@@ -148,13 +149,29 @@ def run_algo(batch_path, uuid, data_path: str = None):
 
     cm.stop_server(dview=dview)
 
-    # Add dictionary to output column of series
-    df.loc[df["uuid"] == uuid, "outputs"] = [d]
-    # Add ran timestamp to ran_time column of series
-    df.loc[df["uuid"] == uuid, "ran_time"] = datetime.now().isoformat(timespec="seconds", sep="T")
-    df.loc[df["uuid"] == uuid, "algo_duration"] = str(round(time.time() - algo_start, 2)) + " sec"
-    # Save DataFrame to disk
-    df.to_pickle(batch_path)
+    # lock batch file while writing back results
+    batch_lock = FileLock(batch_path + '.lock', timeout=30)
+    try:
+        with batch_lock:
+            df = load_batch(batch_path)
+
+            # Add dictionary to output column of series
+            df.loc[df["uuid"] == uuid, "outputs"] = [d]
+            # Add ran timestamp to ran_time column of series
+            df.loc[df["uuid"] == uuid, "ran_time"] = datetime.now().isoformat(timespec="seconds", sep="T")
+            df.loc[df["uuid"] == uuid, "algo_duration"] = str(round(time.time() - algo_start, 2)) + " sec"
+            # Save DataFrame to disk
+            df.to_pickle(batch_path)
+    except Timeout:
+        # Print a message with details in lieu of writing to the batch file
+        msg = f"Batch file could not be written to within {batch_lock.timeout} seconds."
+        if d["success"]:
+            msg += f"\nRun succeeded; results are in {output_dir}."
+        else:
+            msg += f"Run failed. Traceback:\n"
+            msg += d["traceback"]
+
+        raise RuntimeError(msg)
 
 
 @click.command()

requirements.txt

@@ -9,3 +9,4 @@ matplotlib
 click
 psutil
 jupyterlab
+filelock

setup.py

@@ -11,6 +11,7 @@
     "click",
     "psutil",
     "jupyterlab",
+    "filelock"
 ]