From 44c0b20f1fe05f632f12b1e876ad54e1002f0d50 Mon Sep 17 00:00:00 2001
From: rhoadesScholar <rhoadesj@hhmi.org>
Date: Wed, 17 Apr 2024 13:22:24 -0400
Subject: [PATCH] =?UTF-8?q?feat:=20=E2=9C=A8=20Auto=20generate=20rerun=20s?=
 =?UTF-8?q?cript=20for=20relabeling=20step=20of=20segment=20blockwise?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 dacapo/blockwise/scheduler.py | 40 +++++++++++++++++++++++++++++++++++
 dacapo/cli.py                 |  1 -
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/dacapo/blockwise/scheduler.py b/dacapo/blockwise/scheduler.py
index 7dfd72689..4cc2e89bd 100644
--- a/dacapo/blockwise/scheduler.py
+++ b/dacapo/blockwise/scheduler.py
@@ -200,7 +200,47 @@ def segment_blockwise(
     if success:
         shutil.rmtree(tmpdir, ignore_errors=True)
     else:
+        # Write a relabel script to tmpdir
+        output_container = kwargs["output_array_identifier"].container
+        output_dataset = kwargs["output_array_identifier"].dataset
+        out_string = "from dacapo.blockwise import DaCapoBlockwiseTask\n"
+        out_string += (
+            "from dacapo.store.local_array_store import LocalArrayIdentifier\n"
+        )
+        out_string += "import daisy\n"
+        out_string += "from funlib.geometry import Roi, Coordinate\n"
+        out_string += "from upath import UPath as Path\n"
+        out_string += f"output_array_identifier = LocalArrayIdentifier(Path({output_container}), {output_dataset})\n"
+        out_string += (
+            f"total_roi = Roi({total_roi.get_begin()}, {total_roi.get_shape()})\n"
+        )
+        out_string += (
+            f"read_roi = Roi({read_roi.get_begin()}, {read_roi.get_shape()})\n"
+        )
+        out_string += (
+            f"write_roi = Roi({write_roi.get_begin()}, {write_roi.get_shape()})\n"
+        )
+        out_string += "task = DaCapoBlockwiseTask(\n"
+        out_string += f'    "{str(Path(Path(dacapo.blockwise.__file__).parent, "relabel_worker.py"))}"),\n'
+        out_string += "    total_roi,\n"
+        out_string += "    read_roi,\n"
+        out_string += "    write_roi,\n"
+        out_string += f"    {num_workers},\n"
+        out_string += f"    {max_retries},\n"
+        out_string += f"    {timeout},\n"
+        out_string += f"    tmpdir={tmpdir},\n"
+        out_string += f"    output_array_identifier=output_array_identifier,\n"
+        out_string += ")\n"
+        out_string += "success = daisy.run_blockwise([task])\n"
+        out_string += "if success:\n"
+        out_string += f"    shutil.rmtree({tmpdir}, ignore_errors=True)\n"
+        out_string += "else:\n"
+        out_string += '    print("Relabeling failed")\n'
+        with open(Path(tmpdir, "relabel.py"), "w") as f:
+            f.write(out_string)
         raise RuntimeError(
             f"Blockwise segmentation failed. Can rerun with merge files stored at:\n\t{tmpdir}"
+            f"Use read_roi: {read_roi} and write_roi: {write_roi} to rerun."
+            f"Or simply run the script at {Path(tmpdir, 'relabel.py')}"
         )
     return success
diff --git a/dacapo/cli.py b/dacapo/cli.py
index ce34a763f..f4a2a6c41 100644
--- a/dacapo/cli.py
+++ b/dacapo/cli.py
@@ -7,7 +7,6 @@
 import click
 import logging
 from funlib.geometry import Roi, Coordinate
-from funlib.persistence import open_ds
 from dacapo.experiments.datasplits.datasets.dataset import Dataset
 from dacapo.experiments.tasks.post_processors.post_processor_parameters import (
     PostProcessorParameters,