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

Author: akhileshh

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,18 @@
 # pylint: disable=invalid-name, missing-docstring, c-extension-no-member
 
-from datetime import timedelta
-
 import fastremap
 import numpy as np
 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 .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 +27,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,21 +53,22 @@ 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
@@ -76,10 +79,26 @@ def update_chunk(cg: ChunkedGraph, chunk_coords: list[int], layer: int = 2):
     Iterate over all L2 IDs in a chunk and update their cross chunk edges,
     within the periods they were valid/active.
     """
+    global CHILDREN
+
     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)
+
+    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):
+        assert len(CHILDREN) > 0, (nodes, CHILDREN)
+    else:
+        return
+
+    rows = update_nodes(cg, nodes, nodes_ts)
     cg.client.write(rows)

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -15,7 +15,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 +51,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 +63,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 +76,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 +108,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:
@@ -136,7 +137,7 @@ def _update_cross_edges_helper(args):
     parents = cg.get_parents(nodes, fail_to_zero=True)
     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)
@@ -159,7 +160,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)

pychunkedgraph/ingest/upgrade/utils.py

@@ -1,7 +1,7 @@
 # pylint: disable=invalid-name, missing-docstring
 
 from collections import defaultdict
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone
 
 import numpy as np
 from pychunkedgraph.graph import ChunkedGraph
@@ -33,31 +33,50 @@ def get_edit_timestamps(cg: ChunkedGraph, edges_d, start_ts, end_ts) -> list:
     return sorted(timestamps)
 
 
-def get_end_ts(cg: ChunkedGraph, children, start_ts):
-    # get end_ts when node becomes invalid (bigtable resolution is in ms)
-    start = start_ts + timedelta(milliseconds=1)
-    _timestamps = get_parent_timestamps(cg, children, start_time=start)
-    try:
-        end_ts = sorted(_timestamps)[0]
-    except IndexError:
-        # start_ts == end_ts means there has been no edit involving this node
-        # meaning only one timestamp to update cross edges, start_ts
-        end_ts = start_ts
-    return end_ts
+def get_end_timestamps(cg: ChunkedGraph, nodes, nodes_ts, children_map):
+    """
+    Gets the last timestamp for each node at which to update its cross edges.
+    For this, we get parent timestamps for all children of a node.
+    The first timestamp > node_timestamp among these is the last timestamp.
+        This is the timestamp at which one of node's children
+        got a new parent that superseded the current node.
+    """
+    result = []
+    children = np.concatenate([*children_map.values()])
+    timestamps_d = get_parent_timestamps(cg, children)
+    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])
+        try:
+            _timestamps = sorted(_timestamps)
+            _index = np.searchsorted(_timestamps, node_ts)
+            assert _timestamps[_index] == node_ts, (_index, node_ts, _timestamps)
+            end_ts = _timestamps[_index + 1] - timedelta(milliseconds=1)
+        except IndexError:
+            # this node has not been edited, but might have it edges updated
+            end_ts = datetime.now(timezone.utc)
+        result.append(end_ts)
+    return result
 
 
-def get_parent_timestamps(cg: ChunkedGraph, nodes) -> dict[int, set]:
+def get_parent_timestamps(
+    cg: ChunkedGraph, nodes, start_time=None, end_time=None
+) -> dict[int, set]:
     """
     Timestamps of when the given nodes were edited.
     """
+    earliest_ts = cg.get_earliest_timestamp()
     response = cg.client.read_nodes(
         node_ids=nodes,
         properties=[Hierarchy.Parent],
+        start_time=start_time,
+        end_time=end_time,
         end_time_inclusive=False,
     )
 
     result = defaultdict(set)
     for k, v in response.items():
         for cell in v[Hierarchy.Parent]:
-            result[k].add(cell.timestamp)
+            ts = cell.timestamp
+            result[k].add(earliest_ts if ts < earliest_ts else ts)
     return result

pychunkedgraph/utils/redis.py

@@ -18,6 +18,7 @@
 )
 REDIS_PASSWORD = os.environ.get("REDIS_PASSWORD", "")
 REDIS_URL = f"redis://:{REDIS_PASSWORD}@{REDIS_HOST}:{REDIS_PORT}/0"
+CONNECTION = redis.Redis.from_url(REDIS_URL, socket_timeout=60)
 
 keys_fields = ("INGESTION_MANAGER", "JOB_TYPE")
 keys_defaults = ("pcg:imanager", "pcg:job_type")
@@ -27,9 +28,10 @@
 
 
 def get_redis_connection(redis_url=REDIS_URL):
+    if redis_url == REDIS_URL:
+        return CONNECTION
     return redis.Redis.from_url(redis_url, socket_timeout=60)
 
 
 def get_rq_queue(queue):
-    connection = redis.Redis.from_url(REDIS_URL, socket_timeout=60)
-    return Queue(queue, connection=connection)
+    return Queue(queue, connection=CONNECTION)

requirements.in

@@ -10,8 +10,8 @@ google-cloud-datastore>=1.8
 flask
 flask_cors
 python-json-logger
-redis
-rq<2
+redis>5
+rq>2
 pyyaml
 cachetools
 werkzeug

requirements.txt

@@ -97,7 +97,6 @@ gevent==23.9.1
     #   task-queue
 google-api-core[grpc]==2.11.1
     # via
-    #   google-api-core
     #   google-cloud-bigtable
     #   google-cloud-core
     #   google-cloud-datastore
@@ -295,7 +294,7 @@ pytz==2023.3.post1
     # via pandas
 pyyaml==6.0.1
     # via -r requirements.in
-redis==5.0.0
+redis==6.4.0
     # via
     #   -r requirements.in
     #   rq
@@ -316,7 +315,7 @@ rpds-py==0.10.3
     # via
     #   jsonschema
     #   referencing
-rq==1.15.1
+rq==2.4.1
     # via -r requirements.in
 rsa==4.9
     # via