fix(upgrade): use start and end timestamps to filter out irreleveant timestamps

Author: akhileshh

pychunkedgraph/app/meshing/common.py

@@ -4,8 +4,6 @@
 import threading
 
 import numpy as np
-import redis
-from rq import Queue, Connection, Retry
 from flask import Response, current_app, jsonify, make_response, request
 
 from pychunkedgraph import __version__
@@ -145,37 +143,15 @@ def _check_post_options(cg, resp, data, seg_ids):
 def handle_remesh(table_id):
     current_app.request_type = "remesh_enque"
     current_app.table_id = table_id
-    is_priority = request.args.get("priority", True, type=str2bool)
-    is_redisjob = request.args.get("use_redis", False, type=str2bool)
-
     new_lvl2_ids = json.loads(request.data)["new_lvl2_ids"]
-
-    if is_redisjob:
-        with Connection(redis.from_url(current_app.config["REDIS_URL"])):
-
-            if is_priority:
-                retry = Retry(max=3, interval=[1, 10, 60])
-                queue_name = "mesh-chunks"
-            else:
-                retry = Retry(max=3, interval=[60, 60, 60])
-                queue_name = "mesh-chunks-low-priority"
-            q = Queue(queue_name, retry=retry, default_timeout=1200)
-            task = q.enqueue(meshing_tasks.remeshing, table_id, new_lvl2_ids)
-
-        response_object = {"status": "success", "data": {"task_id": task.get_id()}}
-
-        return jsonify(response_object), 202
-    else:
-        new_lvl2_ids = np.array(new_lvl2_ids, dtype=np.uint64)
-        cg = app_utils.get_cg(table_id)
-
-        if len(new_lvl2_ids) > 0:
-            t = threading.Thread(
-                target=_remeshing, args=(cg.get_serialized_info(), new_lvl2_ids)
-            )
-            t.start()
-
-        return Response(status=202)
+    new_lvl2_ids = np.array(new_lvl2_ids, dtype=np.uint64)
+    cg = app_utils.get_cg(table_id)
+    if len(new_lvl2_ids) > 0:
+        t = threading.Thread(
+            target=_remeshing, args=(cg.get_serialized_info(), new_lvl2_ids)
+        )
+        t.start()
+    return Response(status=202)
 
 
 def _remeshing(serialized_cg_info, lvl2_nodes):

pychunkedgraph/graph/edges/__init__.py

@@ -237,11 +237,10 @@ def get_latest_edges(
     Then get supervoxels of those L2 IDs and get parent(s) at `node` level.
     These parents would be the new identities for the stale `partner`.
     """
-    _nodes = np.unique(stale_edges[:, 1])
+    _nodes = np.unique(stale_edges)
     nodes_ts_map = dict(
         zip(_nodes, cg.get_node_timestamps(_nodes, return_numpy=False, normalize=True))
     )
-    _nodes = np.unique(stale_edges)
     layers, coords = cg.get_chunk_layers_and_coordinates(_nodes)
     layers_d = dict(zip(_nodes, layers))
     coords_d = dict(zip(_nodes, coords))
@@ -352,7 +351,9 @@ def _filter(node):
 
         _edges = []
         edges_d = cg.get_cross_chunk_edges(
-            node_ids=l2ids_a, time_stamp=nodes_ts_map[node_b], raw_only=True
+            node_ids=l2ids_a,
+            time_stamp=max(nodes_ts_map[node_a], nodes_ts_map[node_b]),
+            raw_only=True,
         )
         for v in edges_d.values():
             _edges.append(v.get(edge_layer, types.empty_2d))
@@ -382,7 +383,8 @@ def _filter(node):
 
         parents_a = np.array([node_a] * parents_b.size, dtype=basetypes.NODE_ID)
         _new_edges = np.column_stack((parents_a, parents_b))
-        assert _new_edges.size, f"No edge found for {node_a}, {node_b} at {parent_ts}"
+        err = f"No edge found for {node_a}, {node_b} at {edge_layer}; {parent_ts}"
+        assert _new_edges.size, err
         result.append(_new_edges)
     return np.concatenate(result)
 

pychunkedgraph/ingest/upgrade/atomic_layer.py

@@ -1,18 +1,23 @@
 # pylint: disable=invalid-name, missing-docstring, c-extension-no-member
 
-from datetime import timedelta
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import logging, math, time
 
 import fastremap
 import numpy as np
+from tqdm import tqdm
 from pychunkedgraph.graph import ChunkedGraph
-from pychunkedgraph.graph.attributes import Connectivity
+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_parent_timestamps
+from .utils import exists_as_parent, get_end_timestamps, get_parent_timestamps
+
+CHILDREN = {}
 
 
 def update_cross_edges(
-    cg: ChunkedGraph, node, cx_edges_d: dict, node_ts, timestamps: set, earliest_ts
+    cg: ChunkedGraph, node, cx_edges_d: dict, node_ts, node_end_ts, timestamps: set
 ) -> list:
     """
     Helper function to update a single L2 ID.
@@ -27,13 +32,15 @@ def update_cross_edges(
         assert not exists_as_parent(cg, node, edges[:, 0])
         return rows
 
-    partner_parent_ts_d = get_parent_timestamps(cg, 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)
 
     for ts in sorted(timestamps):
-        if ts < earliest_ts:
-            ts = earliest_ts
+        if ts < node_ts:
+            continue
+        if ts > node_end_ts:
+            break
         val_dict = {}
         svs = edges[:, 1]
         parents = cg.get_parents(svs, time_stamp=ts)
@@ -51,35 +58,78 @@ def update_cross_edges(
     return rows
 
 
-def update_nodes(cg: ChunkedGraph, nodes) -> list:
-    nodes_ts = cg.get_node_timestamps(nodes, return_numpy=False, normalize=True)
-    earliest_ts = cg.get_earliest_timestamp()
+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)
     rows = []
-    for node, node_ts in zip(nodes, nodes_ts):
+    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
+            # 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, timestamps_d[node], earliest_ts
+            cg, node, _cx_edges_d, node_ts, end_ts, timestamps_d[node]
         )
         rows.extend(_rows)
     return rows
 
 
-def update_chunk(cg: ChunkedGraph, chunk_coords: list[int], layer: int = 2):
+def _update_nodes_helper(args):
+    cg, nodes, nodes_ts = args
+    return update_nodes(cg, nodes, nodes_ts)
+
+
+def update_chunk(
+    cg: ChunkedGraph, chunk_coords: list[int], layer: int = 2, debug: bool = False
+):
     """
     Iterate over all L2 IDs in a chunk and update their cross chunk edges,
     within the periods they were valid/active.
     """
+    global CHILDREN
+
+    start = time.time()
     x, y, z = chunk_coords
     chunk_id = cg.get_chunk_id(layer=layer, x=x, y=y, z=z)
     cg.copy_fake_edges(chunk_id)
     rr = cg.range_read_chunk(chunk_id)
-    nodes = list(rr.keys())
-    rows = update_nodes(cg, nodes)
-    cg.client.write(rows)
+
+    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)
+
+    if len(nodes) > 0:
+        logging.info(f"Processing {len(nodes)} nodes.")
+        assert len(CHILDREN) > 0, (nodes, CHILDREN)
+    else:
+        return
+
+    if debug:
+        rows = update_nodes(cg, nodes, nodes_ts)
+        cg.client.write(rows)
+    else:
+        task_size = int(math.ceil(len(nodes) / 64))
+        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)
+
+        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())
+
+    print(f"total elaspsed time: {time.time() - start}")

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -3,6 +3,7 @@
 import logging, math, random, time
 import multiprocessing as mp
 from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 import fastremap
 import numpy as np
@@ -15,7 +16,7 @@
 from pychunkedgraph.graph.types import empty_2d
 from pychunkedgraph.utils.general import chunked
 
-from .utils import exists_as_parent, get_parent_timestamps
+from .utils import exists_as_parent, get_end_timestamps, get_parent_timestamps
 
 
 CHILDREN = {}
@@ -51,7 +52,7 @@ def _get_cx_edges_at_timestamp(node, response, ts):
 
 
 def _populate_cx_edges_with_timestamps(
-    cg: ChunkedGraph, layer: int, nodes: list, nodes_ts: list, earliest_ts
+    cg: ChunkedGraph, layer: int, nodes: list, nodes_ts: list
 ):
     """
     Collect timestamps of edits from children, since we use the same timestamp
@@ -63,7 +64,8 @@ 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)
-    for node, node_ts in zip(nodes, nodes_ts):
+    end_timestamps = get_end_timestamps(cg, nodes, nodes_ts, CHILDREN)
+    for node, node_ts, node_end_ts in zip(nodes, nodes_ts, end_timestamps):
         CX_EDGES[node] = {}
         timestamps = timestamps_d[node]
         cx_edges_d_node_ts = _get_cx_edges_at_timestamp(node, response, node_ts)
@@ -75,8 +77,8 @@ def _populate_cx_edges_with_timestamps(
         CX_EDGES[node][node_ts] = cx_edges_d_node_ts
 
         for ts in sorted(timestamps):
-            if ts < earliest_ts:
-                ts = earliest_ts
+            if ts > node_end_ts:
+                break
             CX_EDGES[node][ts] = _get_cx_edges_at_timestamp(node, response, ts)
 
 
@@ -107,7 +109,7 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> l
 
     row_id = serializers.serialize_uint64(node)
     for ts, cx_edges_d in CX_EDGES[node].items():
-        if node_ts > ts:
+        if ts < node_ts:
             continue
         edges = get_latest_edges_wrapper(cg, cx_edges_d, parent_ts=ts)
         if edges.size == 0:
@@ -129,17 +131,29 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> l
     return rows
 
 
+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)
+
+
 def _update_cross_edges_helper(args):
     cg_info, layer, nodes, nodes_ts, earliest_ts = args
     rows = []
     cg = ChunkedGraph(**cg_info)
     parents = cg.get_parents(nodes, fail_to_zero=True)
+
+    tasks = []
     for node, parent, node_ts in zip(nodes, parents, nodes_ts):
         if parent == 0:
-            # invalid id caused by failed ingest task
+            # invalid id caused by failed ingest task / edits
             continue
-        _rows = update_cross_edges(cg, layer, node, node_ts, earliest_ts)
-        rows.extend(_rows)
+        tasks.append((cg, layer, node, node_ts, earliest_ts))
+
+    with ThreadPoolExecutor(max_workers=4) as executor:
+        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)
 
 
@@ -159,7 +173,7 @@ def update_chunk(
     nodes = list(CHILDREN.keys())
     random.shuffle(nodes)
     nodes_ts = cg.get_node_timestamps(nodes, return_numpy=False, normalize=True)
-    _populate_cx_edges_with_timestamps(cg, layer, nodes, nodes_ts, earliest_ts)
+    _populate_cx_edges_with_timestamps(cg, layer, nodes, nodes_ts)
 
     task_size = int(math.ceil(len(nodes) / mp.cpu_count() / 2))
     chunked_nodes = chunked(nodes, task_size)
@@ -171,8 +185,9 @@ def update_chunk(
         args = (cg_info, layer, chunk, ts_chunk, earliest_ts)
         tasks.append(args)
 
-    logging.info(f"Processing {len(nodes)} nodes.")
-    with mp.Pool(min(mp.cpu_count(), len(tasks))) as pool:
+    processes = min(mp.cpu_count() * 2, len(tasks))
+    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),