Remove traces of local_storage

Author: annshress

em_workflows/brt/flow.py

@@ -44,7 +44,6 @@
 from pathlib import Path
 from natsort import os_sorted
 from prefect import task, flow, unmapped
-from prefect.states import Completed, Failed
 from pytools.HedwigZarrImages import HedwigZarrImages
 
 from em_workflows.utils import utils
@@ -649,7 +648,6 @@ def brt_flow(
     # Ref: https://github.com/PrefectHQ/prefect/blob/98d33187ecce032defb8ec7a263de32564e7f7f6/src/prefect/futures.py#L43
     callback_result = list()
     failed = 0
-    total = len(prim_fps)
     for idx, (fp, cb) in enumerate(zip(fps.result(), callback_with_tilt_mov.result())):
         try:
             # In Prefect v3, we can directly get the result without checking state
@@ -659,12 +657,16 @@ def brt_flow(
             callback_result.append(fp.gen_prim_fp_elt(f"Error: {str(e)}"))
             failed += 1
 
-    utils.send_callback_body.submit(
+    send_callback_task = utils.send_callback_body.submit(
         x_no_api=x_no_api,
         token=token,
         callback_url=callback_url,
         files_elts=callback_result,
     )
 
+    utils.final_cleanup_task.submit(
+        fps, x_keep_workdir, wait_for=[utils.allow_failure(send_callback_task)]
+    )
+
     # In Prefect v3, flows should return data, not state objects
     return callback_result

em_workflows/config.py

@@ -107,9 +107,6 @@ class Config:
     user = os.environ["USER"]
     tmp_dir = f"/data/scratch/{user}"
 
-    local_storage = LocalFileSystem(basepath=PREFECT_HOME.value() / "local-storage")
-    pickle_serializer = PickleSerializer(picklelib="pickle")
-
     @staticmethod
     def _mount_point(share_name: str) -> str:
         share = NFS_MOUNT.get(share_name)

em_workflows/czi/flow.py

@@ -273,7 +273,7 @@ async def czi_flow(
     )
 
     # Prefect v3: wait_for is a valid keyword for .submit()
-    cleanup_task = utils.final_cleanup_task.submit(
+    utils.final_cleanup_task.submit(
         fps, x_keep_workdir, wait_for=[utils.allow_failure(send_callback_task)]
     )
 

em_workflows/lrg_2d_rgb/flow.py

@@ -207,22 +207,24 @@ def lrg_2d_flow(
     )
 
     callback_result = list()
-
-    for idx, (fp, cb) in enumerate(zip(fps.result(), callback_with_pyramids)):
-        state = cb.wait()
-        if state.is_completed():
-            callback_result.append(cb.result())
-        else:
-            path = f"{state.state_details.flow_run_id}__{idx}"
-            try:
-                message = LRG2DConfig.local_storage.read_path(path)
-                callback_result.append(fp.gen_prim_fp_elt(message.decode()))
-            except ValueError:
-                callback_result.append(fp.gen_prim_fp_elt("Something went wrong!"))
-
-    utils.send_callback_body.submit(
+    failed = 0
+    for idx, (fp, cb) in enumerate(zip(fps.result(), callback_with_pyramids.result())):
+        try:
+            callback_result.append(cb)
+        except Exception as e:
+            # If there's an error getting the result, use fallback
+            callback_result.append(fp.gen_prim_fp_elt(f"Error: {str(e)}"))
+            failed += 1
+
+    send_callback_task = utils.send_callback_body.submit(
         x_no_api=x_no_api,
         token=token,
         callback_url=callback_url,
         files_elts=callback_result,
     )
+
+    utils.final_cleanup_task.submit(
+        fps, x_keep_workdir, wait_for=[utils.allow_failure(send_callback_task)]
+    )
+
+    return callback_result

em_workflows/sem_tomo/flow.py

@@ -420,21 +420,24 @@ def sem_tomo_flow(
     )
 
     callback_result = list()
-    for idx, (fp, cb) in enumerate(zip(fps.result(), callback_with_corr_movies)):
-        state = cb.wait()
-        if state.is_completed():
-            callback_result.append(cb.result())
-        else:
-            path = f"{state.state_details.flow_run_id}__{idx}"
-            try:
-                message = SEMConfig.local_storage.read_path(path)
-                callback_result.append(fp.gen_prim_fp_elt(message.decode()))
-            except ValueError:
-                callback_result.append(fp.gen_prim_fp_elt("Something went wrong!"))
-
-    utils.send_callback_body.submit(
+    failed = 0
+    for idx, (fp, cb) in enumerate(zip(fps.result(), callback_with_corr_movies.result())):
+        try:
+            callback_result.append(cb)
+        except Exception as e:
+            # If there's an error getting the result, use fallback
+            callback_result.append(fp.gen_prim_fp_elt(f"Error: {str(e)}"))
+            failed += 1
+
+    send_callback_task = utils.send_callback_body.submit(
         x_no_api=x_no_api,
         token=token,
         callback_url=callback_url,
         files_elts=callback_result,
     )
+
+    utils.final_cleanup_task.submit(
+        fps, x_keep_workdir, wait_for=[utils.allow_failure(send_callback_task)]
+    )
+
+    return callback_result

em_workflows/utils/utils.py

@@ -63,28 +63,8 @@ def lookup_dims(fp: Path) -> Header:
         return xyz_cleaned
 
 
-def collect_exception_task_hook(task: Task, task_run: TaskRun, state: State):
-    """
-    This task hook should be used with tasks where you intend to know which step of the flow run broke.
-    Since most of our tasks are mapped by default using filepaths, it takes map index into account as well
-    So that we can notify the user, 'this step of this file broke'.
-    The message is written to a file using prefect's local storage.
-    In order to retrieve it, the flow needs to have a logic at the end,
-        to lookup for this file with exception message if the task run has failed.
-    """
-    message = f"Failure in pipeline step: {task.name}"
-    map_idx = task_run.name.split("-")[-1]
-    flow_run_id = state.state_details.flow_run_id
-    path = f"{flow_run_id}__{map_idx}"
-    try:
-        Config.local_storage.read_path(path)
-    except ValueError:  # ValueError path not found
-        Config.local_storage.write_path(path, message.encode())
-
-
 @task(
     name="mrc to movie generation",
-    on_failure=[collect_exception_task_hook],
 )
 def mrc_to_movie(file_path: FilePath, root: str, asset_type: str, **kwargs):
     """
@@ -357,7 +337,6 @@ def copy_template(working_dir: Path, template_name: str) -> Path:
     name="Batchruntomo conversion",
     tags=["brt"],
     # timeout_seconds=600,
-    on_failure=[collect_exception_task_hook],
 )
 def run_brt(
     file_path: FilePath,
@@ -568,9 +547,6 @@ def notify_api_running(
 #     return ns
 
 
-import asyncio
-
-
 async def notify_api_completion(flow: Flow, flow_run: FlowRun, state: State) -> None:
     """
     https://docs.prefect.io/core/concepts/states.html#overview.
@@ -722,20 +698,6 @@ def send_callback_body(
         )
 
 
-def copy_workdirs_and_cleanup_hook(flow, flow_run, state):
-    stored_result = Config.local_storage.read_path(f"{flow_run.id}__gen_fps")
-    fps: List[FilePath] = Config.pickle_serializer.loads(
-        json.loads(stored_result)["data"].encode()
-    )
-    parameters = flow_run.parameters
-    x_keep_workdir = parameters.get("x_keep_workdir", False)
-
-    for fp in fps:
-        copy_workdir_logs.fn(file_path=fp)
-
-    cleanup_workdir.fn(fps, x_keep_workdir)
-
-
 def callback_with_cleanup(
     fps: List[FilePath],
     callback_result: List,