fix: use registered attributes module

Author: akhileshh

pychunkedgraph/graph/chunkedgraph.py

@@ -672,7 +672,7 @@ def get_subgraph_leaves(
             self, node_id_or_ids, bbox, bbox_is_coordinate, False, True
         )
 
-    def get_edited_edges(
+    def get_edges_from_edits(
         self, chunk_ids: np.ndarray, time_stamp: datetime.datetime = None
     ) -> typing.Dict:
         """
@@ -748,7 +748,7 @@ def get_l2_agglomerations(
         if self.mock_edges is None:
             edges_d = self.read_chunk_edges(chunk_ids)
 
-        edited_edges = self.get_edited_edges(chunk_ids)
+        edited_edges = self.get_edges_from_edits(chunk_ids)
         all_chunk_edges = reduce(
             lambda x, y: x + y,
             chain(edges_d.values(), edited_edges.values()),

pychunkedgraph/graph/edits_sv.py

@@ -5,22 +5,20 @@
 from functools import reduce
 import logging
 import multiprocessing as mp
-from typing import Callable, Iterable
+from typing import Callable
 from datetime import datetime
 from collections import defaultdict, deque
 
 import fastremap
 import numpy as np
 from tqdm import tqdm
-from pychunkedgraph.graph import ChunkedGraph, cache as cache_utils
-from pychunkedgraph.graph.attributes import Connectivity
+from pychunkedgraph.graph import attributes, ChunkedGraph, cache as cache_utils
 from pychunkedgraph.graph.chunks.utils import chunks_overlapping_bbox, get_neighbors
 from pychunkedgraph.graph.cutting_sv import (
     build_kdtrees_by_label,
     pairwise_min_distance_two_sets,
     split_supervoxel_helper,
 )
-from pychunkedgraph.graph.attributes import Hierarchy, OperationLogs
 from pychunkedgraph.graph.edges import Edges
 from pychunkedgraph.graph.types import empty_2d
 from pychunkedgraph.graph import basetypes
@@ -281,9 +279,9 @@ def _add_new_edges(cg: ChunkedGraph, edges_tuple: tuple, time_stamp: datetime =
     for chunk_id in np.unique(chunks):
         val_dict = {}
         mask = chunks_arr[:, 0] == chunk_id
-        val_dict[Connectivity.SplitEdges] = edges[mask]
-        val_dict[Connectivity.Affinity] = affinites[mask]
-        val_dict[Connectivity.Area] = areas[mask]
+        val_dict[attributes.Connectivity.SplitEdges] = edges[mask]
+        val_dict[attributes.Connectivity.Affinity] = affinites[mask]
+        val_dict[attributes.Connectivity.Area] = areas[mask]
         rows.append(
             cg.client.mutate_row(
                 serializers.serialize_uint64(chunk_id, fake_edges=True),
@@ -396,13 +394,15 @@ def copy_parents_and_add_lineage(
     parents = set()
     old_new_map = {k: list(v) for k, v in old_new_map.items()}
     parent_cells_map = cg.client.read_nodes(
-        node_ids=list(old_new_map.keys()), properties=Hierarchy.Parent
+        node_ids=list(old_new_map.keys()), properties=attributes.Hierarchy.Parent
     )
     for old_id, new_ids in old_new_map.items():
         for new_id in new_ids:
             val_dict = {
-                Hierarchy.FormerIdentity: np.array([old_id], dtype=basetypes.NODE_ID),
-                OperationLogs.OperationID: operation_id,
+                attributes.Hierarchy.FormerIdentity: np.array(
+                    [old_id], dtype=basetypes.NODE_ID
+                ),
+                attributes.OperationLogs.OperationID: operation_id,
             }
             result.append(
                 cg.client.mutate_row(serializers.serialize_uint64(new_id), val_dict)
@@ -413,17 +413,19 @@ def copy_parents_and_add_lineage(
                 result.append(
                     cg.client.mutate_row(
                         serializers.serialize_uint64(new_id),
-                        {Hierarchy.Parent: cell.value},
+                        {attributes.Hierarchy.Parent: cell.value},
                         time_stamp=cell.timestamp,
                     )
                 )
-        val_dict = {Hierarchy.NewIdentity: np.array(new_ids, dtype=basetypes.NODE_ID)}
+        val_dict = {
+            attributes.Hierarchy.NewIdentity: np.array(new_ids, dtype=basetypes.NODE_ID)
+        }
         result.append(
             cg.client.mutate_row(serializers.serialize_uint64(old_id), val_dict)
         )
 
     children_cells_map = cg.client.read_nodes(
-        node_ids=list(parents), properties=Hierarchy.Child
+        node_ids=list(parents), properties=attributes.Hierarchy.Child
     )
     for parent, children_cells in children_cells_map.items():
         assert len(children_cells) == 1, children_cells
@@ -437,7 +439,7 @@ def copy_parents_and_add_lineage(
             result.append(
                 cg.client.mutate_row(
                     serializers.serialize_uint64(parent),
-                    {Hierarchy.Child: children},
+                    {attributes.Hierarchy.Child: children},
                     time_stamp=cell.timestamp,
                 )
             )

pychunkedgraph/ingest/cli.py

@@ -5,6 +5,7 @@
 """
 
 import logging
+import os
 
 import click
 import yaml
@@ -70,9 +71,10 @@ def ingest_graph(
     if not retry:
         cg.create()
 
+    get_seg_source_and_destination_ocdbt(cg.meta.data_source.WATERSHED, create=True)
     imanager = IngestionManager(ingest_config, meta)
     enqueue_l2_tasks(imanager, create_atomic_chunk)
-    get_seg_source_and_destination_ocdbt(cg.meta, create=True)
+    os._exit(0)
 
 
 @ingest_cli.command("imanager")

pychunkedgraph/ingest/cluster.py

@@ -143,8 +143,16 @@ def create_atomic_chunk(coords: Sequence[int]):
     for k, v in chunk_edges_active.items():
         logging.debug(f"active_{k}: {len(v)}")
 
-    src, dst = get_seg_source_and_destination_ocdbt(imanager.cg.meta)
-    copy_ws_chunk(imanager.cg, coords, src, dst)
+    src, dst = get_seg_source_and_destination_ocdbt(
+        imanager.cg.meta.data_source.WATERSHED
+    )
+    copy_ws_chunk(
+        src,
+        dst,
+        imanager.cg.meta.graph_config.CHUNK_SIZE,
+        coords,
+        imanager.cg.meta.voxel_bounds,
+    )
     _post_task_completion(imanager, 2, coords)