fix(migration): erase corrupt ids from failed edits

Author: akhileshh

pychunkedgraph/ingest/upgrade/atomic_layer.py

@@ -2,7 +2,7 @@
 
 from collections import defaultdict
 from datetime import datetime, timedelta, timezone
-import logging, time
+import logging, time, os
 from copy import copy
 
 import fastremap
@@ -11,7 +11,7 @@
 from pychunkedgraph.graph.attributes import Connectivity, Hierarchy
 from pychunkedgraph.graph.utils import serializers
 
-from .utils import get_end_timestamps, get_parent_timestamps
+from .utils import fix_corrupt_nodes, get_end_timestamps, get_parent_timestamps
 
 CHILDREN = {}
 
@@ -130,29 +130,37 @@ def update_chunk(cg: ChunkedGraph, chunk_coords: list[int]):
     start = time.time()
     x, y, z = chunk_coords
     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()
+    corrupt_nodes = []
     for k, v in rr.items():
         try:
-            _ = v[Hierarchy.Parent]
-            nodes.append(k)
             CHILDREN[k] = v[Hierarchy.Child][0].value
             ts = v[Hierarchy.Child][0].timestamp
+            _ = v[Hierarchy.Parent]
+            nodes.append(k)
             nodes_ts.append(earliest_ts if ts < earliest_ts else ts)
         except KeyError:
-            # invalid nodes from failed tasks w/o parent column entry
-            continue
+            # ignore invalid nodes from failed ingest tasks, w/o parent column entry
+            # retain invalid nodes from edits to fix the hierarchy
+            if ts > earliest_ts:
+                corrupt_nodes.append(k)
+
+    clean_task = os.environ.get("CLEAN_CHUNKS", "false") == "clean"
+    if clean_task:
+        logging.info(f"found {len(corrupt_nodes)} corrupt nodes {corrupt_nodes[:3]}...")
+        fix_corrupt_nodes(cg, corrupt_nodes, CHILDREN)
+        return
 
-    if len(nodes) > 0:
-        logging.info(f"processing {len(nodes)} nodes.")
-        assert len(CHILDREN) > 0, (nodes, CHILDREN)
-    else:
+    cg.copy_fake_edges(chunk_id)
+    if len(nodes) == 0:
         return
 
+    logging.info(f"processing {len(nodes)} nodes.")
+    assert len(CHILDREN) > 0, (nodes, CHILDREN)
     rows = update_nodes(cg, nodes, nodes_ts)
     cg.client.write(rows)
     logging.info(f"mutations: {len(rows)}, time: {time.time() - start}")

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -1,6 +1,6 @@
 # pylint: disable=invalid-name, missing-docstring, c-extension-no-member
 
-import logging, math, random, time
+import logging, math, random, time, os
 import multiprocessing as mp
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -17,7 +17,7 @@
 from pychunkedgraph.graph.types import empty_2d
 from pychunkedgraph.utils.general import chunked
 
-from .utils import get_end_timestamps, get_parent_timestamps
+from .utils import fix_corrupt_nodes, get_end_timestamps, get_parent_timestamps
 
 
 CHILDREN = {}
@@ -61,6 +61,11 @@ def _populate_cx_edges_with_timestamps(
     when cross edges of children were updated.
     """
 
+    clean_task = os.environ.get("CLEAN_CHUNKS", "false") == "clean"
+    # this data is not needed for clean tasks
+    if clean_task:
+        return
+
     start = time.time()
     global CX_EDGES
     attrs = [Connectivity.CrossChunkEdge[l] for l in range(layer, cg.meta.layer_count)]
@@ -139,20 +144,33 @@ def _update_cross_edges_helper_thread(args):
 
 
 def _update_cross_edges_helper(args):
-    cg_info, layer, nodes, nodes_ts = args
     rows = []
+    clean_task = os.environ.get("CLEAN_CHUNKS", "false") == "clean"
+    cg_info, layer, nodes, nodes_ts = args
     cg = ChunkedGraph(**cg_info)
     parents = cg.get_parents(nodes, fail_to_zero=True)
 
     tasks = []
+    corrupt_nodes = []
+    earliest_ts = cg.get_earliest_timestamp()
     for node, parent, node_ts in zip(nodes, parents, nodes_ts):
         if parent == 0:
-            # invalid id caused by failed ingest task / edits
-            continue
-        tasks.append((cg, layer, node, node_ts))
+            # ignore invalid nodes from failed ingest tasks, w/o parent column entry
+            # retain invalid nodes from edits to fix the hierarchy
+            if node_ts > earliest_ts:
+                corrupt_nodes.append(node)
+        else:
+            tasks.append((cg, layer, node, node_ts))
+
+    if clean_task:
+        logging.info(f"found {len(corrupt_nodes)} corrupt nodes {corrupt_nodes[:3]}...")
+        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]
+        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)
@@ -164,7 +182,7 @@ def update_chunk(
     """
     Iterate over all layer IDs in a chunk and update their cross chunk edges.
     """
-    debug =  nodes is not None
+    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)

pychunkedgraph/ingest/upgrade/utils.py

@@ -1,11 +1,13 @@
 # pylint: disable=invalid-name, missing-docstring
 
 from collections import defaultdict
-from datetime import datetime, timezone
+from datetime import datetime, timedelta
 
 import numpy as np
 from pychunkedgraph.graph import ChunkedGraph
 from pychunkedgraph.graph.attributes import Hierarchy
+from pychunkedgraph.graph.utils import serializers
+from google.cloud.bigtable.row_filters import TimestampRange
 
 
 def exists_as_parent(cg: ChunkedGraph, parent, nodes) -> bool:
@@ -102,3 +104,34 @@ def get_parent_timestamps(
             ts = cell.timestamp
             result[k].add(earliest_ts if ts < earliest_ts else ts)
     return result
+
+
+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.
+    """
+    table = cg.client._table
+    batcher = table.mutations_batcher(flush_count=500)
+    for node in nodes:
+        children = children_d[node]
+        _map = cg.client.read_nodes(node_ids=children, properties=Hierarchy.Parent)
+
+        for child, parent_cells in _map.items():
+            row = table.direct_row(serializers.serialize_uint64(child))
+            for cell in parent_cells:
+                if cell.value == node:
+                    start = cell.timestamp
+                    end = start + timedelta(microseconds=1)
+                    row.delete_cell(
+                        column_family_id=Hierarchy.Parent.family_id,
+                        column=Hierarchy.Parent.key,
+                        time_range=TimestampRange(start=start, end=end),
+                    )
+                    batcher.mutate(row)
+
+        row = table.direct_row(serializers.serialize_uint64(node))
+        row.delete()
+        batcher.mutate(row)
+
+    batcher.flush()