fix: add overwrite guard; avoid range read when subset is passed

CAVEconnectome · Dec 19, 2024 · 7c913f1 · 7c913f1
1 parent 2cdf1f9
commit 7c913f1
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 13 deletions.
diff --git a/pychunkedgraph/meshing/meshgen.py b/pychunkedgraph/meshing/meshgen.py
@@ -955,10 +955,20 @@ def chunk_initial_mesh_task(
     draco_encoding_settings = get_draco_encoding_settings_for_chunk(
         cg, chunk_id, mip, high_padding
     )
+    ids_to_mesh = set()
     if node_id_subset is None:
         seg = get_remapped_segmentation(
             cg, chunk_id, mip, overlap_vx=high_padding, time_stamp=time_stamp
         )
+        try:
+            ts = cg.meta.custom_data["mesh"]["initial_ts"]
+            mesh_ts = datetime.datetime.fromtimestamp(ts)
+        except KeyError:
+            mesh_ts = None
+        range_read = cg.range_read_chunk(
+            chunk_id, properties=attributes.Hierarchy.Child, time_stamp=mesh_ts
+        )
+        ids_to_mesh = set([int(x) for x in range_read.keys()])
     else:
         seg = get_remapped_seg_for_lvl2_nodes(
             cg,
@@ -979,16 +989,6 @@ def chunk_initial_mesh_task(
         print("cv path", mesh_dst)
         print("num ids", len(mesher.ids()))
 
-    try:
-        ts = cg.meta.custom_data["mesh"]["initial_ts"]
-        mesh_ts = datetime.datetime.fromtimestamp(ts)
-    except KeyError:
-        mesh_ts = None
-    range_read = cg.range_read_chunk(
-        chunk_id, properties=attributes.Hierarchy.Child, time_stamp=mesh_ts
-    )
-    ids_to_mesh = set([int(x) for x in range_read.keys()])
-
     result.append(len(mesher.ids()))
     for obj_id in mesher.ids():
         mesh = mesher.get(obj_id, reduction_factor=100, max_error=max_err)
@@ -1009,8 +1009,8 @@ def chunk_initial_mesh_task(
             compress = False
         else:
             file_contents = mesh.to_precomputed()
-            ids_to_mesh.remove(int(obj_id))
             compress = True
+        ids_to_mesh.discard(int(obj_id))
         if WRITING_TO_CLOUD:
             if sharded:
                 merged_meshes[int(obj_id)] = file_contents

diff --git a/pychunkedgraph/meshing/meshing_batch.py b/pychunkedgraph/meshing/meshing_batch.py
@@ -1,8 +1,12 @@
+import argparse, os
+import numpy as np
+from cloudvolume import CloudVolume
+from cloudfiles import CloudFiles
 from taskqueue import TaskQueue, LocalTaskQueue
-import argparse
+
 from pychunkedgraph.graph.chunkedgraph import ChunkedGraph # noqa
-import numpy as np
 from pychunkedgraph.meshing.meshing_sqs import MeshTask
+from pychunkedgraph.meshing import meshgen_utils  # noqa
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -13,11 +17,22 @@
     parser.add_argument('--layer', type=int)
     parser.add_argument('--mip', type=int)
     parser.add_argument('--skip_cache', action='store_true')
+    parser.add_argument('--overwrite', type=bool, default=False)
 
     args = parser.parse_args()
     cache = not args.skip_cache
 
     cg = ChunkedGraph(graph_id=args.cg_name)
+    cv = CloudVolume(
+        f"graphene://https://localhost/segmentation/table/dummy",
+        info=meshgen_utils.get_json_info(cg),
+    )
+    dst = os.path.join(
+        cv.cloudpath, cv.mesh.meta.mesh_path, "initial", str(args.layer)
+    )
+    cf = CloudFiles(dst)
+    if len(list(cf.list())) > 0 and not args.overwrite:
+        raise ValueError(f"Destination {dst} is not empty. Use `--overwrite true` to proceed anyway.")
 
     chunks_arr = []
     for x in range(args.chunk_start[0],args.chunk_end[0]):