fix(edits/migration): pass edge layer to helper function; handle stale source nodes using atomic edges; improve latency with lru cache

Author: akhileshh

pychunkedgraph/graph/edges/__init__.py

@@ -12,6 +12,7 @@
 import tensorstore as ts
 import zstandard as zstd
 from graph_tool import Graph
+from cachetools import LRUCache
 
 from pychunkedgraph.graph import types
 from pychunkedgraph.graph.chunks.utils import (
@@ -21,6 +22,7 @@
 from pychunkedgraph.graph.utils import basetypes
 
 from ..utils import basetypes
+from ..utils.generic import get_parents_at_timestamp
 
 
 _edge_type_fileds = ("in_chunk", "between_chunk", "cross_chunk")
@@ -39,6 +41,7 @@
     ]
 )
 ZSTD_EDGE_COMPRESSION = 17
+PARENTS_CACHE = LRUCache(256 * 1024)
 
 
 class Edges:
@@ -341,7 +344,72 @@ def _filter(node):
         chunks_map[node_b] = np.concatenate(chunks_map[node_b])
         return int(mlayer), _filter(node_a), _filter(node_b)
 
-    def _get_new_edge(edge, parent_ts, padding):
+    def _populate_parents_cache(children: np.ndarray):
+        global PARENTS_CACHE
+
+        not_cached = []
+        for child in children:
+            try:
+                # reset lru index, these will be needed soon
+                _ = PARENTS_CACHE[child]
+            except KeyError:
+                not_cached.append(child)
+
+        all_parents = cg.get_parents(not_cached, current=False)
+        for child, parents in zip(not_cached, all_parents):
+            PARENTS_CACHE[child] = {}
+            for parent, ts in parents:
+                PARENTS_CACHE[child][ts] = parent
+
+    def _get_parents_b(edges, parent_ts, layer):
+        """
+        Attempts to find new partner side nodes.
+        Gets new partners at parent_ts using supervoxels, at `parent_ts`.
+        Searches for new partners that may have any edges to `edges[:,0]`.
+        """
+        children_b = cg.get_children(edges[:, 1], flatten=True)
+        _populate_parents_cache(children_b)
+        _parents_b, missing = get_parents_at_timestamp(
+            children_b, PARENTS_CACHE, time_stamp=parent_ts, unique=True
+        )
+        # handle cache miss cases
+        _parents_b_missing = np.unique(cg.get_parents(missing, time_stamp=parent_ts))
+        parents_b = np.concatenate([_parents_b, _parents_b_missing])
+
+        parents_a = edges[:, 0]
+        stale_a = get_stale_nodes(cg, parents_a, parent_ts=parent_ts)
+        if stale_a.size == parents_a.size:
+            # this is applicable only for v2 to v3 migration
+            # handle cases when source nodes in `edges[:,0]` are stale
+            atomic_edges_d = cg.get_atomic_cross_edges(stale_a)
+            partners = [types.empty_1d]
+            for _edges_d in atomic_edges_d.values():
+                _edges = _edges_d.get(layer, types.empty_2d)
+                partners.append(_edges[:, 1])
+            partners = np.concatenate(partners)
+            return np.unique(cg.get_parents(partners, time_stamp=parent_ts))
+
+        _cx_edges_d = cg.get_cross_chunk_edges(parents_b, time_stamp=parent_ts)
+        _parents_b = []
+        for _node, _edges_d in _cx_edges_d.items():
+            for _edges in _edges_d.values():
+                _mask = np.isin(_edges[:, 1], parents_a)
+                if np.any(_mask):
+                    _parents_b.append(_node)
+        return np.array(_parents_b, dtype=basetypes.NODE_ID)
+
+    def _get_parents_b_with_chunk_mask(
+        l2ids_b: np.ndarray, parents_b: np.ndarray, max_ts: datetime.datetime, edge
+    ):
+        chunks_old = cg.get_chunk_ids_from_node_ids(l2ids_b)
+        chunks_new = cg.get_chunk_ids_from_node_ids(parents_b)
+        chunk_mask = np.isin(chunks_new, chunks_old)
+        parents_b = parents_b[chunk_mask]
+        _stale_nodes = get_stale_nodes(cg, parents_b, parent_ts=max_ts)
+        assert _stale_nodes.size == 0, f"{edge}, {_stale_nodes}, {parent_ts}"
+        return parents_b
+
+    def _get_new_edge(edge, edge_layer, parent_ts, padding):
         """
         Attempts to find new edge(s) for the stale `edge`.
             * Find L2 IDs on opposite sides of the face in L2 chunks along the face.
@@ -353,11 +421,11 @@ def _get_new_edge(edge, parent_ts, padding):
         if l2ids_a.size == 0 or l2ids_b.size == 0:
             return types.empty_2d.copy()
 
-        _edges = []
         max_node_ts = max(nodes_ts_map[node_a], nodes_ts_map[node_b])
         _edges_d = cg.get_cross_chunk_edges(
             node_ids=l2ids_a, time_stamp=max_node_ts, raw_only=True
         )
+        _edges = []
         for v in _edges_d.values():
             if edge_layer in v:
                 _edges.append(v[edge_layer])
@@ -369,27 +437,13 @@ def _get_new_edge(edge, parent_ts, padding):
 
         mask = np.isin(_edges[:, 1], l2ids_b)
         if np.any(mask):
-            parents_a = _edges[mask][:, 0]
-            children_b = cg.get_children(_edges[mask][:, 1], flatten=True)
-            parents_b = np.unique(cg.get_parents(children_b, time_stamp=parent_ts))
-            _cx_edges_d = cg.get_cross_chunk_edges(parents_b, time_stamp=parent_ts)
-            parents_b = []
-            for _node, _edges_d in _cx_edges_d.items():
-                for _edges in _edges_d.values():
-                    _mask = np.isin(_edges[:, 1], parents_a)
-                    if np.any(_mask):
-                        parents_b.append(_node)
-            parents_b = np.array(parents_b, dtype=basetypes.NODE_ID)
+            parents_b = _get_parents_b(_edges[mask], parent_ts, edge_layer)
         else:
             # if none of `l2ids_b` were found in edges, `l2ids_a` already have new edges
             # so get the new identities of `l2ids_b` by using chunk mask
-            parents_b = _edges[:, 1]
-            chunks_old = cg.get_chunk_ids_from_node_ids(l2ids_b)
-            chunks_new = cg.get_chunk_ids_from_node_ids(parents_b)
-            chunk_mask = np.isin(chunks_new, chunks_old)
-            parents_b = parents_b[chunk_mask]
-            _stale_nodes = get_stale_nodes(cg, parents_b, parent_ts=max_node_ts)
-            assert _stale_nodes.size == 0, f"{edge}, {_stale_nodes}, {parent_ts}"
+            parents_b = _get_parents_b_with_chunk_mask(
+                l2ids_b, _edges[:, 1], max_node_ts, edge
+            )
 
         parents_b = np.unique(
             cg.get_roots(parents_b, stop_layer=mlayer, ceil=False, time_stamp=parent_ts)
@@ -402,7 +456,7 @@ def _get_new_edge(edge, parent_ts, padding):
     for edge_layer, _edge in zip(edge_layers, stale_edges):
         max_chebyshev_distance = int(environ.get("MAX_CHEBYSHEV_DISTANCE", 3))
         for pad in range(0, max_chebyshev_distance):
-            _new_edges = _get_new_edge(_edge, parent_ts, padding=pad)
+            _new_edges = _get_new_edge(_edge, edge_layer, parent_ts, padding=pad)
             if _new_edges.size:
                 break
             logging.info(f"{_edge}, expanding search with padding {pad+1}.")
@@ -446,7 +500,7 @@ def get_latest_edges_wrapper(
                 stale_edge_layers,
                 parent_ts=parent_ts,
             )
-            logging.debug(f"{stale_edges} -> {latest_edges}; {parent_ts}")
+            logging.debug(f"{stale_edges} -> {latest_edges[:,1].tolist()}; {parent_ts}")
             _new_cx_edges.append(latest_edges)
         new_cx_edges_d[layer] = np.concatenate(_new_cx_edges)
         nodes.append(np.unique(new_cx_edges_d[layer]))

pychunkedgraph/graph/utils/generic.py

@@ -3,7 +3,6 @@
 TODO categorize properly
 """
 
-
 import datetime
 from typing import Dict
 from typing import Iterable
@@ -173,14 +172,30 @@ def mask_nodes_by_bounding_box(
         adapt_layers = layers - 2
         adapt_layers[adapt_layers < 0] = 0
         fanout = meta.graph_config.FANOUT
-        bounding_box_layer = (
-            bounding_box[None] / (fanout ** adapt_layers)[:, None, None]
-        )
+        bounding_box_layer = bounding_box[None] / (fanout**adapt_layers)[:, None, None]
         bound_check = np.array(
             [
                 np.all(chunk_coordinates < bounding_box_layer[:, 1], axis=1),
                 np.all(chunk_coordinates + 1 > bounding_box_layer[:, 0], axis=1),
             ]
         ).T
 
-        return np.all(bound_check, axis=1)
+        return np.all(bound_check, axis=1)
+
+
+def get_parents_at_timestamp(nodes, parents_ts_map, time_stamp, unique: bool = False):
+    """
+    Search for the first parent with ts <= `time_stamp`.
+    `parents_ts_map[node]` is a map of ts:parent with sorted timestamps (desc).
+    """
+    skipped_nodes = []
+    parents = set() if unique else []
+    for node in nodes:
+        try:
+            for ts, parent in parents_ts_map[node].items():
+                if time_stamp >= ts:
+                    parents.add(parent) if unique else parents.append(parent)
+                    break
+        except KeyError:
+            skipped_nodes.append(node)
+    return list(parents), skipped_nodes

pychunkedgraph/ingest/cluster.py

@@ -10,7 +10,7 @@
 from typing import Callable, Dict, Iterable, Tuple, Sequence
 
 import numpy as np
-from rq import Queue as RQueue
+from rq import Queue as RQueue, Retry
 
 
 from .utils import chunk_id_str, get_chunks_not_done, randomize_grid_points
@@ -209,6 +209,7 @@ def _queue_tasks(imanager: IngestionManager, chunk_fn: Callable, coords: Iterabl
                     timeout=environ.get("L2JOB_TIMEOUT", "3m"),
                     result_ttl=0,
                     job_id=chunk_id_str(2, chunk_coord),
+                    retry=Retry(int(environ.get("RETRY_COUNT", 1))),
                 )
             )
         q.enqueue_many(job_datas)

pychunkedgraph/ingest/upgrade/atomic_layer.py

@@ -10,26 +10,13 @@
 from pychunkedgraph.graph import ChunkedGraph, types
 from pychunkedgraph.graph.attributes import Connectivity, Hierarchy
 from pychunkedgraph.graph.utils import serializers
+from pychunkedgraph.graph.utils.generic import get_parents_at_timestamp
 
 from .utils import fix_corrupt_nodes, get_end_timestamps, get_parent_timestamps
 
 CHILDREN = {}
 
 
-def _get_parents_at_timestamp(nodes, parents_ts_map, time_stamp):
-    """
-    Search for the first parent with ts <= `time_stamp`.
-    `parents_ts_map[node]` is a map of ts:parent with sorted timestamps (desc).
-    """
-    parents = []
-    for node in nodes:
-        for ts, parent in parents_ts_map[node].items():
-            if time_stamp >= ts:
-                parents.append(parent)
-                break
-    return parents
-
-
 def update_cross_edges(
     cg: ChunkedGraph,
     node,
@@ -59,7 +46,7 @@ def update_cross_edges(
             break
 
         val_dict = {}
-        parents = _get_parents_at_timestamp(partners, parents_ts_map, ts)
+        parents, _ = get_parents_at_timestamp(partners, parents_ts_map, ts)
         edge_parents_d = dict(zip(partners, parents))
         for layer, layer_edges in cx_edges_d.items():
             layer_edges = fastremap.remap(

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -10,9 +10,8 @@
 import numpy as np
 from tqdm import tqdm
 
-from pychunkedgraph.graph import ChunkedGraph
+from pychunkedgraph.graph import ChunkedGraph, edges
 from pychunkedgraph.graph.attributes import Connectivity, Hierarchy
-from pychunkedgraph.graph.edges import get_latest_edges_wrapper
 from pychunkedgraph.graph.utils import serializers
 from pychunkedgraph.graph.types import empty_2d
 from pychunkedgraph.utils.general import chunked
@@ -105,7 +104,6 @@ def _populate_cx_edges_with_timestamps(
             row_id = serializers.serialize_uint64(node)
             val_dict = {Hierarchy.StaleTimeStamp: 0}
             rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp=node_end_ts))
-
     cg.client.write(rows)
 
 
@@ -119,7 +117,7 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts) -> list:
     for ts, cx_edges_d in CX_EDGES[node].items():
         if ts < node_ts:
             continue
-        cx_edges_d, edge_nodes = get_latest_edges_wrapper(cg, cx_edges_d, parent_ts=ts)
+        cx_edges_d, edge_nodes = edges.get_latest_edges_wrapper(cg, cx_edges_d, parent_ts=ts)
         if edge_nodes.size == 0:
             continue
 
@@ -138,13 +136,7 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts) -> list:
     return rows
 
 
-def _update_cross_edges_helper_thread(args):
-    cg, layer, node, node_ts = args
-    return update_cross_edges(cg, layer, node, node_ts)
-
-
 def _update_cross_edges_helper(args):
-    rows = []
     clean_task = os.environ.get("CLEAN_CHUNKS", "false") == "clean"
     cg_info, layer, nodes, nodes_ts = args
     cg = ChunkedGraph(**cg_info)
@@ -167,12 +159,9 @@ def _update_cross_edges_helper(args):
         fix_corrupt_nodes(cg, corrupt_nodes, CHILDREN)
         return
 
-    with ThreadPoolExecutor(max_workers=4) as executor:
-        futures = [
-            executor.submit(_update_cross_edges_helper_thread, task) for task in tasks
-        ]
-        for future in tqdm(as_completed(futures), total=len(futures)):
-            rows.extend(future.result())
+    rows = []
+    for task in tasks:
+        rows.extend(update_cross_edges(*task))
     cg.client.write(rows)
 
 
@@ -204,12 +193,13 @@ def update_chunk(
 
     if debug:
         rows = []
+        logging.info(f"processing {len(nodes)} nodes with 1 worker.")
         for node, node_ts in zip(nodes, nodes_ts):
             rows.extend(update_cross_edges(cg, layer, node, node_ts))
         logging.info(f"total elaspsed time: {time.time() - start}")
         return
 
-    task_size = int(math.ceil(len(nodes) / mp.cpu_count() / 2))
+    task_size = int(math.ceil(len(nodes) / mp.cpu_count()))
     chunked_nodes = chunked(nodes, task_size)
     chunked_nodes_ts = chunked(nodes_ts, task_size)
     cg_info = cg.get_serialized_info()

pychunkedgraph/ingest/upgrade/utils.py

@@ -108,8 +108,8 @@ def get_parent_timestamps(
 
 def fix_corrupt_nodes(cg: ChunkedGraph, nodes: list, children_d: dict):
     """
-    Iteratively removes a node from parent column of its children.
-    Then removes the node iteself, effectively erasing it.
+    For each node: delete it from parent column of its children.
+    Then deletes the node itself, effectively erasing it from hierarchy.
     """
     table = cg.client._table
     batcher = table.mutations_batcher(flush_count=500)

pychunkedgraph/ingest/utils.py

@@ -9,7 +9,7 @@
 
 import numpy as np
 import tensorstore as ts
-from rq import Queue, Worker
+from rq import Queue, Retry, Worker
 from rq.worker import WorkerStatus
 
 from . import IngestConfig
@@ -199,6 +199,7 @@ def queue_layer_helper(parent_layer: int, imanager: IngestionManager, fn):
                     result_ttl=0,
                     job_id=chunk_id_str(parent_layer, chunk_coord),
                     timeout=f"{timeout_scale * int(parent_layer * parent_layer)}m",
+                    retry=Retry(int(environ.get("RETRY_COUNT", 1))),
                 )
             )
         q.enqueue_many(job_datas)