feat(upgrade): cache stale timestamp info; remove unnecessary checks to reduce latency

Author: akhileshh

pychunkedgraph/graph/attributes.py

@@ -160,6 +160,12 @@ class Hierarchy:
         serializer=serializers.NumPyValue(dtype=basetypes.NODE_ID),
     )
 
+    # track when nodes became stale, required for migration
+    # will be eventually deleted by GC rule for column family_id 3.
+    StaleTimeStamp = _Attribute(
+        key=b"stale_ts", family_id="3", serializer=serializers.Pickle()
+    )
+
 
 class GraphMeta:
     key = b"meta"

pychunkedgraph/graph/edges/__init__.py

@@ -201,22 +201,23 @@ def get_edges(source: str, nodes: np.ndarray) -> Edges:
 
 
 def get_stale_nodes(
-    cg, edge_nodes: Iterable[basetypes.NODE_ID], parent_ts: datetime.datetime = None
+    cg, nodes: Iterable[basetypes.NODE_ID], parent_ts: datetime.datetime = None
 ):
     """
-    Checks to see if partner nodes in edges (edges[:,1]) are stale.
-    This is done by getting a supervoxel of the node and check
+    Checks to see if given nodes are stale.
+    This is done by getting a supervoxel of a node and checking
     if it has a new parent at the same layer as the node.
     """
-    edge_supervoxels = cg.get_single_leaf_multiple(edge_nodes)
+    nodes = np.array(nodes, dtype=basetypes.NODE_ID)
+    supervoxels = cg.get_single_leaf_multiple(nodes)
     # nodes can be at different layers due to skip connections
-    edge_nodes_layers = cg.get_chunk_layers(edge_nodes)
+    node_layers = cg.get_chunk_layers(nodes)
     stale_nodes = [types.empty_1d]
-    for layer in np.unique(edge_nodes_layers):
-        _mask = edge_nodes_layers == layer
-        layer_nodes = edge_nodes[_mask]
+    for layer in np.unique(node_layers):
+        _mask = node_layers == layer
+        layer_nodes = nodes[_mask]
         _nodes = cg.get_roots(
-            edge_supervoxels[_mask],
+            supervoxels[_mask],
             stop_layer=layer,
             ceil=False,
             time_stamp=parent_ts,

pychunkedgraph/ingest/cluster.py

@@ -135,7 +135,7 @@ def upgrade_atomic_chunk(coords: Sequence[int]):
     redis = get_redis_connection()
     imanager = IngestionManager.from_pickle(redis.get(r_keys.INGESTION_MANAGER))
     coords = np.array(list(coords), dtype=int)
-    update_atomic_chunk(imanager.cg, coords, layer=2)
+    update_atomic_chunk(imanager.cg, coords)
     _post_task_completion(imanager, 2, coords)
 
 

pychunkedgraph/ingest/upgrade/atomic_layer.py

@@ -1,50 +1,53 @@
 # pylint: disable=invalid-name, missing-docstring, c-extension-no-member
 
+from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import timedelta
 import logging, math, time
+from copy import copy
 
 import fastremap
 import numpy as np
 from tqdm import tqdm
-from pychunkedgraph.graph import ChunkedGraph
+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 exists_as_parent, get_end_timestamps, get_parent_timestamps
+from .utils import get_end_timestamps, get_parent_timestamps
 
 CHILDREN = {}
 
 
 def update_cross_edges(
-    cg: ChunkedGraph, node, cx_edges_d: dict, node_ts, node_end_ts, timestamps: set
+    cg: ChunkedGraph,
+    node,
+    cx_edges_d: dict,
+    node_ts,
+    node_end_ts,
+    timestamps_d: defaultdict[int, set],
 ) -> list:
     """
     Helper function to update a single L2 ID.
     Returns a list of mutations with given timestamps.
     """
     rows = []
     edges = np.concatenate(list(cx_edges_d.values()))
-    uparents = np.unique(cg.get_parents(edges[:, 0], time_stamp=node_ts))
-    assert uparents.size <= 1, f"{node}, {node_ts}, {uparents}"
-    if uparents.size == 0 or node != uparents[0]:
-        # if node is not the parent at this ts, it must be invalid
-        assert not exists_as_parent(cg, node, edges[:, 0])
-        return rows
+    partners = np.unique(edges[:, 1])
 
-    partner_parent_ts_d = get_parent_timestamps(cg, np.unique(edges[:, 1]))
-    for v in partner_parent_ts_d.values():
-        timestamps.update(v)
+    timestamps = copy(timestamps_d[node])
+    for partner in partners:
+        timestamps.update(timestamps_d[partner])
 
     for ts in sorted(timestamps):
         if ts < node_ts:
             continue
         if ts > node_end_ts:
             break
+
         val_dict = {}
-        svs = edges[:, 1]
-        parents = cg.get_parents(svs, time_stamp=ts)
-        edge_parents_d = dict(zip(svs, parents))
+        parents = cg.get_parents(partners, time_stamp=ts)
+        edge_parents_d = dict(zip(partners, parents))
         for layer, layer_edges in cx_edges_d.items():
             layer_edges = fastremap.remap(
                 layer_edges, edge_parents_d, preserve_missing_labels=True
@@ -61,20 +64,26 @@ def update_cross_edges(
 def update_nodes(cg: ChunkedGraph, nodes, nodes_ts, children_map=None) -> list:
     if children_map is None:
         children_map = CHILDREN
-    end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, children_map)
-    timestamps_d = get_parent_timestamps(cg, nodes)
+    end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, children_map, layer=2)
+
     cx_edges_d = cg.get_atomic_cross_edges(nodes)
+    all_cx_edges = [types.empty_2d]
+    for _cx_edges_d in cx_edges_d.values():
+        if _cx_edges_d:
+            all_cx_edges.append(np.concatenate(list(_cx_edges_d.values())))
+    all_partners = np.unique(np.concatenate(all_cx_edges)[:, 1])
+    timestamps_d = get_parent_timestamps(cg, np.concatenate([nodes, all_partners]))
+
     rows = []
     for node, node_ts, end_ts in zip(nodes, nodes_ts, end_timestamps):
-        if cg.get_parent(node) is None:
-            # invalid id caused by failed ingest task / edits
-            continue
+        end_ts -= timedelta(milliseconds=1)
         _cx_edges_d = cx_edges_d.get(node, {})
         if not _cx_edges_d:
             continue
-        _rows = update_cross_edges(
-            cg, node, _cx_edges_d, node_ts, end_ts, timestamps_d[node]
-        )
+        _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))
         rows.extend(_rows)
     return rows
 
@@ -84,9 +93,7 @@ def _update_nodes_helper(args):
     return update_nodes(cg, nodes, nodes_ts)
 
 
-def update_chunk(
-    cg: ChunkedGraph, chunk_coords: list[int], layer: int = 2, debug: bool = False
-):
+def update_chunk(cg: ChunkedGraph, chunk_coords: list[int], debug: bool = False):
     """
     Iterate over all L2 IDs in a chunk and update their cross chunk edges,
     within the periods they were valid/active.
@@ -95,35 +102,41 @@ 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)
+    chunk_id = cg.get_chunk_id(layer=2, x=x, y=y, z=z)
     cg.copy_fake_edges(chunk_id)
     rr = cg.range_read_chunk(chunk_id)
 
     nodes = []
     nodes_ts = []
     earliest_ts = cg.get_earliest_timestamp()
     for k, v in rr.items():
-        nodes.append(k)
-        CHILDREN[k] = v[Hierarchy.Child][0].value
-        ts = v[Hierarchy.Child][0].timestamp
-        nodes_ts.append(earliest_ts if ts < earliest_ts else ts)
+        try:
+            _ = v[Hierarchy.Parent]
+            nodes.append(k)
+            CHILDREN[k] = v[Hierarchy.Child][0].value
+            ts = v[Hierarchy.Child][0].timestamp
+            nodes_ts.append(earliest_ts if ts < earliest_ts else ts)
+        except KeyError:
+            # invalid nodes from failed tasks w/o parent column entry
+            continue
 
     if len(nodes) > 0:
-        logging.info(f"Processing {len(nodes)} nodes.")
+        logging.info(f"processing {len(nodes)} nodes.")
         assert len(CHILDREN) > 0, (nodes, CHILDREN)
     else:
         return
 
     if debug:
         rows = update_nodes(cg, nodes, nodes_ts)
     else:
-        task_size = int(math.ceil(len(nodes) / 64))
+        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:
@@ -132,4 +145,4 @@ def update_chunk(
                 rows.extend(future.result())
 
     cg.client.write(rows)
-    print(f"total elaspsed time: {time.time() - start}")
+    logging.info(f"total elaspsed time: {time.time() - start}")

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -16,7 +16,7 @@
 from pychunkedgraph.graph.types import empty_2d
 from pychunkedgraph.utils.general import chunked
 
-from .utils import exists_as_parent, get_end_timestamps, get_parent_timestamps
+from .utils import get_end_timestamps, get_parent_timestamps
 
 
 CHILDREN = {}
@@ -64,7 +64,9 @@ def _populate_cx_edges_with_timestamps(
     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)
+    end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, CHILDREN, layer=layer)
+
+    rows = []
     for node, node_ts, node_end_ts in zip(nodes, nodes_ts, end_timestamps):
         CX_EDGES[node] = {}
         timestamps = timestamps_d[node]
@@ -81,32 +83,18 @@ def _populate_cx_edges_with_timestamps(
                 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))
+    cg.client.write(rows)
+
 
-def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> list:
+def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts) -> list:
     """
     Helper function to update a single ID.
     Returns a list of mutations with timestamps.
     """
     rows = []
-    if node_ts > earliest_ts:
-        try:
-            cx_edges_d = CX_EDGES[node][node_ts]
-        except KeyError:
-            raise KeyError(f"{node}:{node_ts}")
-        edges = np.concatenate([empty_2d] + list(cx_edges_d.values()))
-        if edges.size:
-            parents = cg.get_roots(
-                edges[:, 0], time_stamp=node_ts, stop_layer=layer, ceil=False
-            )
-            uparents = np.unique(parents)
-            layers = cg.get_chunk_layers(uparents)
-            uparents = uparents[layers == layer]
-            assert uparents.size <= 1, f"{node}, {node_ts}, {uparents}"
-            if uparents.size == 0 or node != uparents[0]:
-                # if node is not the parent at this ts, it must be invalid
-                assert not exists_as_parent(cg, node, edges[:, 0]), f"{node}, {node_ts}"
-                return rows
-
     row_id = serializers.serialize_uint64(node)
     for ts, cx_edges_d in CX_EDGES[node].items():
         if ts < node_ts:
@@ -132,12 +120,12 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> l
 
 
 def _update_cross_edges_helper_thread(args):
-    cg, layer, node, node_ts, earliest_ts = args
-    return update_cross_edges(cg, layer, node, node_ts, earliest_ts)
+    cg, layer, node, node_ts = args
+    return update_cross_edges(cg, layer, node, node_ts)
 
 
 def _update_cross_edges_helper(args):
-    cg_info, layer, nodes, nodes_ts, earliest_ts = args
+    cg_info, layer, nodes, nodes_ts = args
     rows = []
     cg = ChunkedGraph(**cg_info)
     parents = cg.get_parents(nodes, fail_to_zero=True)
@@ -147,7 +135,7 @@ def _update_cross_edges_helper(args):
         if parent == 0:
             # invalid id caused by failed ingest task / edits
             continue
-        tasks.append((cg, layer, node, node_ts, earliest_ts))
+        tasks.append((cg, layer, node, node_ts))
 
     with ThreadPoolExecutor(max_workers=4) as executor:
         futures = [executor.submit(_update_cross_edges_helper_thread, task) for task in tasks]
@@ -163,10 +151,10 @@ def update_chunk(
     """
     Iterate over all layer IDs in a chunk and update their cross chunk edges.
     """
+    debug =  nodes is not None
     start = time.time()
     x, y, z = chunk_coords
     chunk_id = cg.get_chunk_id(layer=layer, x=x, y=y, z=z)
-    earliest_ts = cg.get_earliest_timestamp()
     _populate_nodes_and_children(cg, chunk_id, nodes=nodes)
     if not CHILDREN:
         return
@@ -175,23 +163,31 @@ def update_chunk(
     nodes_ts = cg.get_node_timestamps(nodes, return_numpy=False, normalize=True)
     _populate_cx_edges_with_timestamps(cg, layer, nodes, nodes_ts)
 
+    if debug:
+        rows = []
+        for node, node_ts in zip(nodes, nodes_ts):
+            rows.extend(update_cross_edges(cg, layer, node, node_ts))
+        cg.client.write(rows)
+        logging.info(f"total elaspsed time: {time.time() - start}")
+        return
+
     task_size = int(math.ceil(len(nodes) / mp.cpu_count() / 2))
     chunked_nodes = chunked(nodes, task_size)
     chunked_nodes_ts = chunked(nodes_ts, task_size)
     cg_info = cg.get_serialized_info()
 
     tasks = []
     for chunk, ts_chunk in zip(chunked_nodes, chunked_nodes_ts):
-        args = (cg_info, layer, chunk, ts_chunk, earliest_ts)
+        args = (cg_info, layer, chunk, ts_chunk)
         tasks.append(args)
 
     processes = min(mp.cpu_count() * 2, len(tasks))
-    logging.info(f"Processing {len(nodes)} nodes with {processes} workers.")
+    logging.info(f"processing {len(nodes)} nodes with {processes} workers.")
     with mp.Pool(processes) as pool:
         _ = list(
             tqdm(
                 pool.imap_unordered(_update_cross_edges_helper, tasks),
                 total=len(tasks),
             )
         )
-    print(f"total elaspsed time: {time.time() - start}")
+    logging.info(f"total elaspsed time: {time.time() - start}")