perf(upgrade): reduce latency for atomic layer chunks

Author: akhileshh

pychunkedgraph/ingest/upgrade/atomic_layer.py

@@ -1,31 +1,43 @@
 # pylint: disable=invalid-name, missing-docstring, c-extension-no-member
 
 from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timedelta, timezone
-import logging, math, time
+import logging, time
 from copy import copy
 
 import fastremap
 import numpy as np
-from tqdm import tqdm
 from pychunkedgraph.graph import ChunkedGraph, types
 from pychunkedgraph.graph.attributes import Connectivity, Hierarchy
 from pychunkedgraph.graph.utils import serializers
-from pychunkedgraph.utils.general import chunked
 
 from .utils import 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,
     cx_edges_d: dict,
     node_ts,
     node_end_ts,
-    timestamps_d: defaultdict[int, set],
+    timestamps_map: defaultdict[int, set],
+    parents_ts_map: defaultdict[int, dict],
 ) -> list:
     """
     Helper function to update a single L2 ID.
@@ -35,9 +47,9 @@ def update_cross_edges(
     edges = np.concatenate(list(cx_edges_d.values()))
     partners = np.unique(edges[:, 1])
 
-    timestamps = copy(timestamps_d[node])
+    timestamps = copy(timestamps_map[node])
     for partner in partners:
-        timestamps.update(timestamps_d[partner])
+        timestamps.update(timestamps_map[partner])
 
     node_end_ts = node_end_ts or datetime.now(timezone.utc)
     for ts in sorted(timestamps):
@@ -47,7 +59,7 @@ def update_cross_edges(
             break
 
         val_dict = {}
-        parents = cg.get_parents(partners, time_stamp=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(
@@ -63,6 +75,7 @@ def update_cross_edges(
 
 
 def update_nodes(cg: ChunkedGraph, nodes, nodes_ts, children_map=None) -> list:
+    start = time.time()
     if children_map is None:
         children_map = CHILDREN
     end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, children_map, layer=2)
@@ -75,31 +88,39 @@ def update_nodes(cg: ChunkedGraph, nodes, nodes_ts, children_map=None) -> list:
     all_partners = np.unique(np.concatenate(all_cx_edges)[:, 1])
     timestamps_d = get_parent_timestamps(cg, np.concatenate([nodes, all_partners]))
 
+    parents_ts_map = defaultdict(dict)
+    all_parents = cg.get_parents(all_partners, current=False)
+    for partner, parents in zip(all_partners, all_parents):
+        for parent, ts in parents:
+            parents_ts_map[partner][ts] = parent
+    logging.info(f"update_nodes init {len(nodes)}: {time.time() - start}")
+
     rows = []
+    skipped = []
     for node, node_ts, end_ts in zip(nodes, nodes_ts, end_timestamps):
         is_stale = end_ts is not None
         _cx_edges_d = cx_edges_d.get(node, {})
-        if not _cx_edges_d:
-            continue
         if is_stale:
             end_ts -= timedelta(milliseconds=1)
-
-        _rows = update_cross_edges(cg, node, _cx_edges_d, node_ts, end_ts, timestamps_d)
-        if is_stale:
             row_id = serializers.serialize_uint64(node)
             val_dict = {Hierarchy.StaleTimeStamp: 0}
-            _rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp=end_ts))
-        rows.extend(_rows)
-
-    return rows
+            rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp=end_ts))
 
+        if not _cx_edges_d:
+            skipped.append(node)
+            continue
 
-def _update_nodes_helper(args):
-    cg, nodes, nodes_ts = args
-    return update_nodes(cg, nodes, nodes_ts)
+        _rows = update_cross_edges(
+            cg, node, _cx_edges_d, node_ts, end_ts, timestamps_d, parents_ts_map
+        )
+        rows.extend(_rows)
+    parents = cg.get_roots(skipped)
+    layers = cg.get_chunk_layers(parents)
+    assert np.all(layers == cg.meta.layer_count)
+    return rows
 
 
-def update_chunk(cg: ChunkedGraph, chunk_coords: list[int], debug: bool = False):
+def update_chunk(cg: ChunkedGraph, chunk_coords: list[int]):
     """
     Iterate over all L2 IDs in a chunk and update their cross chunk edges,
     within the periods they were valid/active.
@@ -132,23 +153,6 @@ def update_chunk(cg: ChunkedGraph, chunk_coords: list[int], debug: bool = False)
     else:
         return
 
-    if debug:
-        rows = update_nodes(cg, nodes, nodes_ts)
-    else:
-        task_size = int(math.ceil(len(nodes) / 16))
-        chunked_nodes = chunked(nodes, task_size)
-        chunked_nodes_ts = chunked(nodes_ts, task_size)
-        tasks = []
-        for chunk, ts_chunk in zip(chunked_nodes, chunked_nodes_ts):
-            args = (cg, chunk, ts_chunk)
-            tasks.append(args)
-        logging.info(f"task size {task_size}, count {len(tasks)}.")
-
-        rows = []
-        with ThreadPoolExecutor(max_workers=8) as executor:
-            futures = [executor.submit(_update_nodes_helper, task) for task in tasks]
-            for future in tqdm(as_completed(futures), total=len(futures)):
-                rows.extend(future.result())
-
+    rows = update_nodes(cg, nodes, nodes_ts)
     cg.client.write(rows)
-    logging.info(f"total elaspsed time: {time.time() - start}")
+    logging.info(f"mutations: {len(rows)}, time: {time.time() - start}")