perf(migration): cache parent information with timestamps to reduce io latency

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,6 +344,23 @@ def _filter(node):
         chunks_map[node_b] = np.concatenate(chunks_map[node_b])
         return int(mlayer), _filter(node_a), _filter(node_b)
 
+    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_new_edge(edge, parent_ts, padding):
         """
         Attempts to find new edge(s) for the stale `edge`.
@@ -371,7 +391,13 @@ def _get_new_edge(edge, parent_ts, padding):
         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))
+            # parents_b = np.unique(cg.get_parents(children_b, time_stamp=parent_ts))
+            _populate_parents_cache(children_b)
+            _parents_b, missing = get_parents_at_timestamp(
+                children_b, PARENTS_CACHE, time_stamp=parent_ts, unique=True
+            )
+            _parents_b_missing = np.unique(cg.get_parents(missing, time_stamp=parent_ts))
+            parents_b = np.concatenate([_parents_b, _parents_b_missing])
             _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():

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(2),
                 )
             )
         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(2),
                 )
             )
         q.enqueue_many(job_datas)