fix(edits): batch create higher layer ids, reuse future roots from lock, sanity check defaults to true

Author: akhileshh

pychunkedgraph/graph/cache.py

@@ -179,6 +179,7 @@ def parents_multiple(
             time_stamp=time_stamp,
             fail_to_zero=fail_to_zero,
         )
+        mask = mask | (parents == 0)
         update(self.parents_cache, node_ids[~mask], parents[~mask])
         return parents
 

pychunkedgraph/graph/chunkedgraph.py

@@ -214,7 +214,7 @@ def get_parents(
                 end_time=time_stamp,
                 end_time_inclusive=True,
             )
-            if not parent_rows:
+            if not parent_rows and not fail_to_zero:
                 return types.empty_1d
 
             parents = []
@@ -736,8 +736,8 @@ def get_l2_agglomerations(
             else:
                 all_chunk_edges = all_chunk_edges.get_pairs()
             supervoxels = self.get_children(level2_ids, flatten=True)
-            mask0 = np.in1d(all_chunk_edges[:, 0], supervoxels)
-            mask1 = np.in1d(all_chunk_edges[:, 1], supervoxels)
+            mask0 = np.isin(all_chunk_edges[:, 0], supervoxels)
+            mask1 = np.isin(all_chunk_edges[:, 1], supervoxels)
             return all_chunk_edges[mask0 & mask1]
 
         l2id_children_d = self.get_children(level2_ids)
@@ -809,7 +809,7 @@ def add_edges(
         source_coords: typing.Sequence[int] = None,
         sink_coords: typing.Sequence[int] = None,
         allow_same_segment_merge: typing.Optional[bool] = False,
-        do_sanity_check: typing.Optional[bool] = False,
+        do_sanity_check: typing.Optional[bool] = True,
     ) -> operation.GraphEditOperation.Result:
         """
         Adds an edge to the chunkedgraph
@@ -842,7 +842,7 @@ def remove_edges(
         path_augment: bool = True,
         disallow_isolating_cut: bool = True,
         bb_offset: typing.Tuple[int, int, int] = (240, 240, 24),
-        do_sanity_check: typing.Optional[bool] = False,
+        do_sanity_check: typing.Optional[bool] = True,
     ) -> operation.GraphEditOperation.Result:
         """
         Removes edges - either directly or after applying a mincut

pychunkedgraph/graph/client/bigtable/client.py

@@ -467,7 +467,7 @@ def lock_roots(
 
             lock_results = {}
             root_ids = np.unique(new_root_ids)
-            max_workers = max(1, len(root_ids) // 2)
+            max_workers = max(1, len(root_ids))
             with ThreadPoolExecutor(max_workers=max_workers) as executor:
                 future_to_root = {
                     executor.submit(self.lock_root, root_id, operation_id): root_id
@@ -518,7 +518,7 @@ def lock_roots_indefinitely(
 
         root_ids = np.unique(new_root_ids)
         lock_results = {}
-        max_workers = max(1, len(root_ids) // 2)
+        max_workers = max(1, len(root_ids))
         failed_to_lock = []
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
             future_to_root = {
@@ -601,7 +601,7 @@ def renew_lock(self, root_id: np.uint64, operation_id: np.uint64) -> bool:
 
     def renew_locks(self, root_ids: np.uint64, operation_id: np.uint64) -> bool:
         """Renews existing root node locks with operation_id to extend time."""
-        max_workers = max(1, len(root_ids) // 2)
+        max_workers = max(1, len(root_ids))
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
             futures = {
                 executor.submit(self.renew_lock, root_id, operation_id): root_id
@@ -640,7 +640,7 @@ def get_consolidated_lock_timestamp(
         """Minimum of multiple lock timestamps."""
         if len(root_ids) == 0:
             return None
-        max_workers = max(1, len(root_ids) // 2)
+        max_workers = max(1, len(root_ids))
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
             futures = {
                 executor.submit(self.get_lock_timestamp, root_id, op_id): (

pychunkedgraph/graph/edges/__init__.py

@@ -207,7 +207,7 @@ def flip_ids(id_map, node_ids):
     return np.concatenate(ids).astype(basetypes.NODE_ID)
 
 
-def _get_new_nodes(
+def get_new_nodes(
     cg, nodes: np.ndarray, layer: int, parent_ts: datetime.datetime = None
 ):
     unique_nodes, inverse = np.unique(nodes, return_inverse=True)
@@ -248,7 +248,7 @@ def get_stale_nodes(
     for layer in np.unique(node_layers):
         _mask = node_layers == layer
         layer_nodes = nodes[_mask]
-        _nodes = _get_new_nodes(cg, supervoxels[_mask], layer, parent_ts)
+        _nodes = get_new_nodes(cg, supervoxels[_mask], layer, parent_ts)
         stale_mask = layer_nodes != _nodes
         stale_nodes.append(layer_nodes[stale_mask])
     return np.concatenate(stale_nodes)
@@ -447,8 +447,11 @@ def _get_parents_b(edges, parent_ts, layer, fallback: bool = False):
         Searches for new partners that may have any edges to `edges[:,0]`.
         """
         if PARENTS_CACHE is None:
+            # this cache is set only during migration
+            # also, fallback is not applicable if no migration
             children_b = cg.get_children(edges[:, 1], flatten=True)
             parents_b = np.unique(cg.get_parents(children_b, time_stamp=parent_ts))
+            fallback = False
         else:
             children_b = _get_children_from_cache(edges[:, 1])
             _populate_parents_cache(children_b)
@@ -564,7 +567,7 @@ def _get_new_edge(edge, edge_layer, parent_ts, padding, fallback: bool = False):
                     if fallback:
                         parents_b = _get_parents_b(_edges, parent_ts, edge_layer, True)
 
-        parents_b = np.unique(_get_new_nodes(cg, parents_b, mlayer, parent_ts))
+        parents_b = np.unique(get_new_nodes(cg, parents_b, mlayer, parent_ts))
         parents_a = np.array([node_a] * parents_b.size, dtype=basetypes.NODE_ID)
         return np.column_stack((parents_a, parents_b))
 

pychunkedgraph/graph/edits.py

@@ -1,6 +1,6 @@
 # pylint: disable=invalid-name, missing-docstring, too-many-locals, c-extension-no-member
 
-import datetime, random
+import datetime, logging, random
 from typing import Dict
 from typing import List
 from typing import Tuple
@@ -16,7 +16,7 @@
 from . import types
 from . import attributes
 from . import cache as cache_utils
-from .edges import get_latest_edges_wrapper, flip_ids
+from .edges import get_latest_edges_wrapper, flip_ids, get_new_nodes
 from .edges.utils import concatenate_cross_edge_dicts
 from .edges.utils import merge_cross_edge_dicts
 from .utils import basetypes
@@ -25,6 +25,8 @@
 from ..utils.general import in2d
 from ..debug.utils import sanity_check, sanity_check_single
 
+logger = logging.getLogger(__name__)
+
 
 def _init_old_hierarchy(cg, l2ids: np.ndarray, parent_ts: datetime.datetime = None):
     """
@@ -684,11 +686,8 @@ def _update_cross_edge_cache_batched(self, new_ids: list):
 
         # Distribute results back to each parent's cache
         # Key insight: edges[:, 0] are children, map them to their parent
-        edge_parents = self.cg.get_roots(
-            edge_nodes,
-            stop_layer=parent_layer,
-            ceil=False,
-            time_stamp=self._last_ts,
+        edge_parents = get_new_nodes(
+            self.cg, edge_nodes, parent_layer, self._last_ts
         )
         edge_parents_d = dict(zip(edge_nodes, edge_parents))
         for new_id in new_ids:
@@ -714,6 +713,48 @@ def _update_cross_edge_cache_batched(self, new_ids: list):
             self.cg.cache.cross_chunk_edges_cache[new_id] = parent_cx_edges_d
         return updated_entries
 
+    def _get_new_ids(self, chunk_id, count, is_root):
+        batch_size = count
+        new_ids = []
+        while len(new_ids) < count:
+            candidate_ids = self.cg.id_client.create_node_ids(
+                chunk_id, batch_size, root_chunk=is_root
+            )
+            existing = self.cg.client.read_nodes(node_ids=candidate_ids)
+            non_existing = set(candidate_ids) - existing.keys()
+            new_ids.extend(non_existing)
+            batch_size = min(batch_size * 2, 2**16)
+        return new_ids[:count]
+
+    def _get_new_parents(self, layer, ccs, graph_ids) -> tuple[dict, dict]:
+        cc_layer_chunk_map = {}
+        size_map = defaultdict(int)
+        for i, cc_idx in enumerate(ccs):
+            parent_layer = layer + 1  # must be reset for each connected component
+            cc_ids = graph_ids[cc_idx]
+            if len(cc_ids) == 1:
+                # skip connection
+                parent_layer = self.cg.meta.layer_count
+                cx_edges_d = self.cg.get_cross_chunk_edges(
+                    [cc_ids[0]], time_stamp=self._last_ts
+                )
+                for l in range(layer + 1, self.cg.meta.layer_count):
+                    if len(cx_edges_d[cc_ids[0]].get(l, types.empty_2d)) > 0:
+                        parent_layer = l
+                        break
+            chunk_id = self.cg.get_parent_chunk_id(cc_ids[0], parent_layer)
+            cc_layer_chunk_map[i] = (parent_layer, chunk_id)
+            size_map[chunk_id] += 1
+
+        chunk_ids = list(size_map.keys())
+        random.shuffle(chunk_ids)
+        chunk_new_ids_map = {}
+        layers = self.cg.get_chunk_layers(chunk_ids)
+        for c, l in zip(chunk_ids, layers):
+            is_root = l == self.cg.meta.layer_count
+            chunk_new_ids_map[c] = self._get_new_ids(c, size_map[c], is_root)
+        return chunk_new_ids_map, cc_layer_chunk_map
+
     def _create_new_parents(self, layer: int):
         """
         keep track of old IDs
@@ -726,37 +767,13 @@ def _create_new_parents(self, layer: int):
         """
         new_ids = self._new_ids_d[layer]
         layer_node_ids = self._get_layer_node_ids(new_ids, layer)
-        components, graph_ids = self._get_connected_components(layer_node_ids, layer)
-        for cc_indices in components:
-            parent_layer = layer + 1  # must be reset for each connected component
-            cc_ids = graph_ids[cc_indices]
-            if len(cc_ids) == 1:
-                # skip connection
-                parent_layer = self.cg.meta.layer_count
-                cx_edges_d = self.cg.get_cross_chunk_edges(
-                    [cc_ids[0]], time_stamp=self._last_ts
-                )
-                for l in range(layer + 1, self.cg.meta.layer_count):
-                    if len(cx_edges_d[cc_ids[0]].get(l, types.empty_2d)) > 0:
-                        parent_layer = l
-                        break
+        ccs, _ids = self._get_connected_components(layer_node_ids, layer)
+        new_parents_map, cc_layer_chunk_map = self._get_new_parents(layer, ccs, _ids)
 
-            # TODO: handle skip connected root id creation separately
-            chunk_id = self.cg.get_parent_chunk_id(cc_ids[0], parent_layer)
-            is_root = parent_layer == self.cg.meta.layer_count
-            batch_size = 1
-            parent = None
-            while parent is None:
-                candidate_ids = self.cg.id_client.create_node_ids(
-                    chunk_id, batch_size, root_chunk=is_root
-                )
-                existing = self.cg.client.read_nodes(node_ids=candidate_ids)
-                for cid in candidate_ids:
-                    if cid not in existing:
-                        parent = cid
-                        break
-                if parent is None:
-                    batch_size = min(batch_size * 2, 2**16)
+        for i, cc_indices in enumerate(ccs):
+            cc_ids = _ids[cc_indices]
+            parent_layer, chunk_id = cc_layer_chunk_map[i]
+            parent = new_parents_map[chunk_id].pop()
 
             self._new_ids_d[parent_layer].append(parent)
             self._update_id_lineage(parent, cc_ids, layer, parent_layer)
@@ -786,19 +803,20 @@ def run(self) -> Iterable:
         """
         self._new_ids_d[2] = self._new_l2_ids
         for layer in range(2, self.cg.meta.layer_count):
-            if len(self._new_ids_d[layer]) == 0:
+            new_nodes = self._new_ids_d[layer]
+            if len(new_nodes) == 0:
                 continue
-            self.cg.cache.new_ids.update(self._new_ids_d[layer])
+            self.cg.cache.new_ids.update(new_nodes)
             # all new IDs in this layer have been created
             # update their cross chunk edges and their neighbors'
             with self._profiler.profile(f"l{layer}_update_cx_cache"):
-                entries = self._update_cross_edge_cache_batched(self._new_ids_d[layer])
+                entries = self._update_cross_edge_cache_batched(new_nodes)
                 self.new_entries.extend(entries)
 
             with self._profiler.profile(f"l{layer}_update_neighbor_cx"):
                 entries = _update_neighbor_cx_edges(
                     self.cg,
-                    self._new_ids_d[layer],
+                    new_nodes,
                     self._new_old_id_d,
                     self._old_new_id_d,
                     time_stamp=self._time_stamp,
@@ -861,10 +879,24 @@ def _get_cross_edges_val_dicts(self):
         return val_dicts
 
     def create_new_entries(self) -> List:
+        max_layer = self.cg.meta.layer_count
         val_dicts = self._get_cross_edges_val_dicts()
-        for layer in range(2, self.cg.meta.layer_count + 1):
+        for layer in range(2, max_layer + 1):
             new_ids = self._new_ids_d[layer]
             for id_ in new_ids:
+                if self.do_sanity_check:
+                    root_layer = self.cg.get_chunk_layer(self.cg.get_root(id_))
+                    assert root_layer == max_layer, (id_, self.cg.get_root(id_))
+
+                    if layer < max_layer:
+                        try:
+                            _parent = self.cg.get_parent(id_)
+                            _children = self.cg.get_children(_parent)
+                            assert id_ in _children, (layer, id_, _parent, _children)
+                        except TypeError as e:
+                            logger.error(id_, _parent, self.cg.get_root(id_))
+                            raise TypeError from e
+
                 val_dict = val_dicts.get(id_, {})
                 children = self.cg.get_children(id_)
                 err = f"parent layer less than children; op {self._opid}"