fix(migration): use None for non stale node end ts, performance improvements

Author: akhileshh

pychunkedgraph/ingest/upgrade/atomic_layer.py

@@ -2,7 +2,7 @@
 
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone
 import logging, math, time
 from copy import copy
 
@@ -39,6 +39,7 @@ def update_cross_edges(
     for partner in partners:
         timestamps.update(timestamps_d[partner])
 
+    node_end_ts = node_end_ts or datetime.now(timezone.utc)
     for ts in sorted(timestamps):
         if ts < node_ts:
             continue
@@ -76,15 +77,20 @@ def update_nodes(cg: ChunkedGraph, nodes, nodes_ts, children_map=None) -> list:
 
     rows = []
     for node, node_ts, end_ts in zip(nodes, nodes_ts, end_timestamps):
-        end_ts -= timedelta(milliseconds=1)
+        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)
-        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))
+        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
 
 

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -4,6 +4,7 @@
 import multiprocessing as mp
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime, timezone
 
 import fastremap
 import numpy as np
@@ -59,33 +60,47 @@ def _populate_cx_edges_with_timestamps(
     for all IDs involved in an edit, we can use the timestamps of
     when cross edges of children were updated.
     """
+
+    start = time.time()
     global CX_EDGES
     attrs = [Connectivity.CrossChunkEdge[l] for l in range(layer, cg.meta.layer_count)]
     all_children = np.concatenate(list(CHILDREN.values()))
     response = cg.client.read_nodes(node_ids=all_children, properties=attrs)
     timestamps_d = get_parent_timestamps(cg, nodes)
     end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, CHILDREN, layer=layer)
+    logging.info(f"_populate_nodes_and_children init: {time.time() - start}")
 
-    rows = []
-    for node, node_ts, node_end_ts in zip(nodes, nodes_ts, end_timestamps):
+    start = time.time()
+    partners_map = {}
+    for node, node_ts in zip(nodes, nodes_ts):
         CX_EDGES[node] = {}
-        timestamps = timestamps_d[node]
         cx_edges_d_node_ts = _get_cx_edges_at_timestamp(node, response, node_ts)
-
         edges = np.concatenate([empty_2d] + list(cx_edges_d_node_ts.values()))
-        partner_parent_ts_d = get_parent_timestamps(cg, edges[:, 1])
-        for v in partner_parent_ts_d.values():
-            timestamps.update(v)
+        partners_map[node] = edges[:, 1]
         CX_EDGES[node][node_ts] = cx_edges_d_node_ts
 
+    partners = np.unique(np.concatenate([*partners_map.values()]))
+    partner_parent_ts_d = get_parent_timestamps(cg, partners)
+    logging.info(f"get partners timestamps init: {time.time() - start}")
+
+    rows = []
+    for node, node_ts, node_end_ts in zip(nodes, nodes_ts, end_timestamps):
+        timestamps = timestamps_d[node]
+        for partner in partners_map[node]:
+            timestamps.update(partner_parent_ts_d[partner])
+
+        is_stale = node_end_ts is not None
+        node_end_ts = node_end_ts or datetime.now(timezone.utc)
         for ts in sorted(timestamps):
             if ts > node_end_ts:
                 break
             CX_EDGES[node][ts] = _get_cx_edges_at_timestamp(node, response, ts)
 
-        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))
+        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=node_end_ts))
+
     cg.client.write(rows)
 
 
@@ -140,7 +155,6 @@ def _update_cross_edges_helper(args):
         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())
-
     cg.client.write(rows)
 
 
@@ -154,13 +168,21 @@ def update_chunk(
     start = time.time()
     x, y, z = chunk_coords
     chunk_id = cg.get_chunk_id(layer=layer, x=x, y=y, z=z)
+
     _populate_nodes_and_children(cg, chunk_id, nodes=nodes)
+    logging.info(f"_populate_nodes_and_children: {time.time() - start}")
     if not CHILDREN:
         return
     nodes = list(CHILDREN.keys())
     random.shuffle(nodes)
+
+    start = time.time()
     nodes_ts = cg.get_node_timestamps(nodes, return_numpy=False, normalize=True)
+    logging.info(f"get_node_timestamps: {time.time() - start}")
+
+    start = time.time()
     _populate_cx_edges_with_timestamps(cg, layer, nodes, nodes_ts)
+    logging.info(f"_populate_cx_edges_with_timestamps: {time.time() - start}")
 
     if debug:
         rows = []

pychunkedgraph/ingest/upgrade/utils.py

@@ -64,16 +64,19 @@ def get_end_timestamps(
 
     for node, node_ts in zip(nodes, nodes_ts):
         node_children = children_map[node]
-        _timestamps = set().union(*[timestamps_d[k] for k in node_children])
+        _children_timestamps = []
+        for k in node_children:
+            if k in timestamps_d:
+                _children_timestamps.append(timestamps_d[k])
+        _timestamps = set().union(*_children_timestamps)
         _timestamps.add(node_ts)
         try:
             _timestamps = sorted(_timestamps)
             _index = np.searchsorted(_timestamps, node_ts)
-            assert _timestamps[_index] == node_ts, (_index, node_ts, _timestamps)
             end_ts = _timestamps[_index + 1]
         except IndexError:
             # this node has not been edited, but might have it edges updated
-            end_ts = datetime.now(timezone.utc)
+            end_ts = None
         result.append(end_ts)
     return result