fix(upgrade): remove unnecessary checks to reduce latency

Author: akhileshh

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,52 @@
 # pylint: disable=invalid-name, missing-docstring, c-extension-no-member
 
+from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 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
@@ -62,19 +64,21 @@ 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)
+
     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
         _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)
         rows.extend(_rows)
     return rows
 
@@ -84,9 +88,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,7 +97,7 @@ 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)
 
@@ -108,22 +110,29 @@ def update_chunk(
         ts = v[Hierarchy.Child][0].timestamp
         nodes_ts.append(earliest_ts if ts < earliest_ts else ts)
 
+    nodes = np.array(nodes, dtype=np.uint64)
+    node_parents = cg.get_parents(nodes, fail_to_zero=True)
+    invalid_mask = node_parents != 0
+    nodes = nodes[invalid_mask]
+    nodes_ts = nodes_ts[invalid_mask]
+
     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 +141,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

@@ -88,25 +88,6 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> l
     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:
@@ -186,7 +167,7 @@ def update_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(

pychunkedgraph/ingest/upgrade/utils.py

@@ -61,7 +61,7 @@ def get_end_timestamps(cg: ChunkedGraph, nodes, nodes_ts, children_map):
 
 def get_parent_timestamps(
     cg: ChunkedGraph, nodes, start_time=None, end_time=None
-) -> dict[int, set]:
+) -> defaultdict[int, set]:
     """
     Timestamps of when the given nodes were edited.
     """

pychunkedgraph/utils/general.py

@@ -26,10 +26,6 @@ def reverse_dictionary(dictionary):
 
 
 def chunked(l: Sequence, n: int):
-    """
-    Yield successive n-sized chunks from l.
-    NOTE: Use itertools.batched from python 3.12
-    """
     """
     Yield successive n-sized chunks from l.
     NOTE: Use itertools.batched from python 3.12
@@ -39,9 +35,6 @@ def chunked(l: Sequence, n: int):
     it = iter(l)
     while batch := tuple(islice(it, n)):
         yield batch
-    it = iter(l)
-    while batch := tuple(islice(it, n)):
-        yield batch
 
 
 def in2d(arr1: np.ndarray, arr2: np.ndarray) -> np.ndarray: